Numpy.where() подробное, пошаговое объяснение/примеры

Ответ 1

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

Интуитивно понятно, что np.where напоминает запрос "скажите, где в этом массиве записи удовлетворяют заданному условию".

>>> a = np.arange(5,10)
>>> np.where(a < 8)       # tell me where in a, entries are < 8
(array([0, 1, 2]),)       # answer: entries indexed by 0, 1, 2

Его также можно использовать для получения записей в массиве, которые удовлетворяют условию:

>>> a[np.where(a < 8)] 
array([5, 6, 7])          # selects from a entries 0, 1, 2

Когда a является двумерным массивом, np.where() возвращает массив идентификаторов строк и массив идентификаторов строк:

>>> a = np.arange(4,10).reshape(2,3)
array([[4, 5, 6],
       [7, 8, 9]])
>>> np.where(a > 8)
(array(1), array(2))

Как и в случае с 1d, мы можем использовать np.where() для получения записей в массиве 2d, которые удовлетворяют условию:

>>> a[np.where(a > 8)] # selects from a entries 0, 1, 2

массив ([9])

Обратите внимание, что когда a равен 1d, np.where() по-прежнему возвращает массив строки idx и массив col idx, но столбцы имеют длину 1, поэтому последний является пустым массивом.

Ответ 2

Вот немного веселее. Я обнаружил, что очень часто NumPy делает именно то, что я бы хотел, чтобы это делало - иногда мне быстрее просто попробовать что-то, чем читать документы. На самом деле лучше всего использовать смесь обоих.

Я думаю, что ваш ответ в порядке (и это нормально, если хотите). Это просто "лишний".

import numpy as np

a = np.arange(4,10).reshape(2,3)

wh = np.where(a>7)
gt = a>7
x  = np.where(gt)

print "wh: ", wh
print "gt: ", gt
print "x:  ", x

дает:

wh:  (array([1, 1]), array([1, 2]))
gt:  [[False False False]
      [False  True  True]]
x:   (array([1, 1]), array([1, 2]))

... но:

print "a[wh]: ", a[wh]
print "a[gt]  ", a[gt]
print "a[x]:  ", a[x]

дает:

a[wh]:  [8 9]
a[gt]   [8 9]
a[x]:   [8 9]