Выберите первые элементы, которые удовлетворяют условию в двумерном массиве

Question

У меня есть массив с формой [N, M, 2]. У меня есть условие для последнего индекса, т.е.

condition = matrix[:,:,0]**2 + matrix[:,:,1]**2 <1

Теперь я хотел бы выбрать элементы матрицы, которые удовлетворяют этому условию. Если я слепо применяю matrix[condition], я, очевидно, получаю результат с формой [что-то, 2], потому что иначе это не будет матрица.

Моя проблема в том, что Мне нужно сохранить форму [N, ..., 2]. Например, было бы хорошо взять только первые K элементов, которые удовлетворяют условию (при условии, что всегда есть хотя бы K). Как я могу это сделать?

С помощью для l oop я бы написал

matrix_new = np.zeros([N,K,2])
for i in range(N):
    condition = matrix[i,:,0]**2 + matrix[i,:,1]**2 <1
    matrix_new[i] = matrix[condition][:K]

, но это было бы очень медленно. Как я могу избежать для l oop?

Answer 1

Вы уже вырезали свое новое пространство. Вам просто нужно сохранить в нем свою условную матрицу:

matrix           = np.array([[[1,2],[3,4]],[[3,4],[5,6]],[[7,8],[0,1]]])
cond             = matrix[...,0]**2+matrix[...,1]**2 < 10 
matrix_new       = np.zeros([3,2,2])
matrix_new[cond] = matrix[cond]

Помните, что массив [cond] возвращает представление в numpy перед результатом. Следовательно, numpy имеет несколько трюков в рукаве и приведёт к тому, что вы сохраните, прежде чем вернуть упрощенный массив. : -)

РЕДАКТИРОВАТЬ:

Единственное, что я мог найти, это np.take_along_axis и методы np.argsort для получения определенных векторов в определенном порядке на основе условной матрицы. Примерно так:

matrix     = np.array([[[1,2],[3,4]],[[1,1],[5,6]],[[7,8],[0,1]]])
cond       = matrix[...,0]**2 + matrix[...,1]**2
matrix_new = np.take_along_axis(matrix,cond.argsort()[...,None],axis=1)

matrix_new[:,0,:] #Smallest vectors in each. 
matrix_new[:,1,:] #Second smallest vectors in each. 

Кроме того, вы можете получить значение K с помощью «K = max ((cond <10) .sum (1))». </p>