Использование numpy.argpartition игнорируя NaN - PullRequest
Новая
       15

Использование numpy.argpartition игнорируя NaN

2 голосов
/ 30 мая 2019

У меня есть большой массив с примерно 49 миллионами элементов (7000 * 7000), и там мне нужно найти самые большие N элементов и их индексы, игнорирующие все NaN. Я не могу удалить эти NaN заранее, потому что мне нужны значения индекса самых больших N элементов из первого массива, чтобы извлечь данные из другого, у которого есть NaN в разных индексах по сравнению с первым массивом. Я пытался

np.argpartition(first_array, -N)[-N:]

это работает очень хорошо для массива без NaN, но если есть NaN, то Nan будут самым большим элементом, потому что в питоне это считается бесконечностью. 

x = np.array([np.nan, 2, -1, 2, -4, -8, -9, 6, -3]).reshape(3, 3)
y = np.argpartition(x.ravel() , -3)[-3:]
z = x.ravel()[y]
# this is the result I am getting  === [2, 6, nan]
# but I need this ==== [2, 2, 6]

1 Ответ

1 голос
/ 30 мая 2019

Использовать количество NaN для смещения и, таким образом, вычисления индексов и извлечения значений - 

In [200]: N = 3

In [201]: c = np.isnan(x).sum()

In [204]: idx = np.argpartition(x.ravel() , -N-c)[-N-c:-c]

In [207]: val = x.flat[idx]

In [208]: idx,val
Out[208]: (array([1, 3, 7]), array([2., 2., 6.]))
...