Найти индексы элементов, которые соответствуют условию 2D массива по строке

Question

У меня есть двумерный массив с поплавками, который называется A с размерами n x m. И у меня есть одномерный массив с плавающей точкой измерения n, называемый B. Для каждой строки i в A я хочу найти индексы значений, которые меньше, чем i-е значение в B

Я пытался использовать np.argwhere

A = np.array([[.1,.2,.3,.4], [.1,.2,.3,.4], [.1,.2,.3,.4]]) B = np.array([0.15, 0.25, 0.35]) np.argwhere(A.T < B) Это вывод

array([[0, 0],[0, 1],[0, 2],[1, 1],[1, 2],[2, 2]])

Я бы хотел, чтобы это было что-то вроде

[array([0,1,2]), array([1,2]), array([2])]

И, конечно, я бы хотел, чтобы оно было векторизовано, а не зациклено.

Answer 1

Использование np.where + np.split

x, y = np.where(A.T < B)

s = np.flatnonzero(np.append([False], x[1:] != x[:-1]))

np.split(y, s)

[array([0, 1, 2], dtype=int64),
 array([1, 2], dtype=int64),
 array([2], dtype=int64)]