Function to sort array
def sort012 ( a , arr_size ) : NEW_LINE INDENT lo = 0 NEW_LINE hi = arr_size - 1 NEW_LINE mid = 0 NEW_LINE while mid <= hi : NEW_LINE INDENT if a [ mid ] == 0 : NEW_LINE INDENT a [ lo ] , a [ mid ] = a [ mid ] , a [ lo ] NEW_LINE lo = lo + 1 NEW_LINE mid = mid + 1 NEW_LINE DEDENT elif a [ mid ] == 1 : NEW_LINE INDENT mid = mid + 1 NEW_LINE DEDENT else : NEW_LINE INDENT a [ mid ] , a [ hi ] = a [ hi ] , a [ mid ] NEW_LINE hi = hi - 1 NEW_LINE DEDENT DEDENT return a NEW_LINE DEDENT