Numpy выберите строки в зависимости от условия

Question

Я пытаюсь использовать numpy для удаления строк из двумерного массива, где первое значение строки (поэтому элемент с индексом 0) не соответствует определенному условию.

Я могусделать это с обычным питоном, используя два цикла, но я хотел бы сделать это более эффективно с NumPy, например, с numpy.where

Я пробовал разные вещи с numpy.where и numpy.delete, но я борюсь стот факт, что я хочу выбирать строки, используя условие, которое нужно проверять только первым элементом, а не вторым (меня не интересует значение второго элемента)

Вот пример, гдеЯ хочу сохранить только те строки, в которых первое значение каждой строки равно 6.

Ввод:

[[0,4],
[0,5],
[3,5],
[6,8],
[9,1],
[6,1]]

Вывод:

[[6,8],
[6,1]]

Answer 1

Используйте логическую маску:

mask = z[:, 0] == 6
z[mask, :]

Это намного эффективнее, чем np.where, потому что вы можете использовать логическую маску напрямую, без необходимости сначала преобразовывать ее в массив индексов.

Один вкладыш:

z[z[:,0]==6, :]

Answer 2

Программа:

import numpy as np
np_array = np.array([[0,4],[0,5],[3,5],[6,8],[9,1],[6,1]])
rows=np.where(np_array[:,0]==6)
print(np_array[rows])

Вывод:

[[6 8]
 [6 1]]

И если вы хотите попасть в 2d список, используйте

np_array[rows].tolist()

Вывод 2d списка

[[6, 8], [6, 1]]