Есть ли более быстрая альтернатива np.where для определения индексов?

Question

У меня есть такой массив:

arrayElements = [[1, 4, 6],[2, 4, 6],[3, 5, 6],...,[2, 5, 6]]

Мне нужно знать, например, индексы, где arrayElements равен 1.

Right сейчас я делаю:

rows, columns = np.where(arrayElements == 1)

Это работает, но я делаю это в al oop, который перебирает все возможные значения элемента, в моем случае это 1-500000 +. Это займет 30-40 минут, в зависимости от размера моего массива. Кто-нибудь может предложить лучший способ сделать это? (Дополнительная информация заключается в том, что мне не важен столбец, в котором находится значение, только строка, я не уверен, что это полезно.)

Редактировать: мне нужно знать значение каждого элемента в отдельности. То есть мне нужны значения строк для каждого значения, которое содержит элементы.

Answer 1

Итак, вы генерируете тысячи массивов, таких как:

In [271]: [(i,np.where(arr==i)[0]) for i in range(1,7)]                                                
Out[271]: 
[(1, array([0])),
 (2, array([1, 3])),
 (3, array([2])),
 (4, array([0, 1])),
 (5, array([2, 3])),
 (6, array([0, 1, 2, 3]))]

Я мог бы провести тест == для всех значений одновременно с небольшим количеством широковещания:

In [281]: arr==np.arange(1,7)[:,None,None]                                                             
Out[281]: 
array([[[ True, False, False],
        [False, False, False],
        [False, False, False],
        [False, False, False]],

       [[False, False, False],
        [ True, False, False],
        [False, False, False],
        [ True, False, False]],

       [[False, False, False],
        [False, False, False],
        [ True, False, False],
        [False, False, False]],

       [[False,  True, False],
        [False,  True, False],
        [False, False, False],
        [False, False, False]],

       [[False, False, False],
        [False, False, False],
        [False,  True, False],
        [False,  True, False]],

       [[False, False,  True],
        [False, False,  True],
        [False, False,  True],
        [False, False,  True]]])

и так как вы заботитесь только о строках, примените any:

In [282]: (arr==np.arange(1,7)[:,None,None]).any(axis=2)                                               
Out[282]: 
array([[ True, False, False, False],
       [False,  True, False,  True],
       [False, False,  True, False],
       [ True,  True, False, False],
       [False, False,  True,  True],
       [ True,  True,  True,  True]])

. where для этого значения совпадает со значениями в Out [271], но сгруппированы по-разному:

In [283]: np.where((arr==np.arange(1,7)[:,None,None]).any(axis=2))                                     
Out[283]: 
(array([0, 1, 1, 2, 3, 3, 4, 4, 5, 5, 5, 5]),
 array([0, 1, 3, 2, 0, 1, 2, 3, 0, 1, 2, 3]))

Его можно разделить на:

In [284]: from collections import defaultdict                                                          
In [285]: dd = defaultdict(list)                                                                       
In [287]: for i,j in zip(*Out[283]): dd[i].append(j)                                                   
In [288]: dd                                                                                           
Out[288]: 
defaultdict(list,
            {0: [0], 1: [1, 3], 2: [2], 3: [0, 1], 4: [2, 3], 5: [0, 1, 2, 3]})

Этот 2-й подход может быть более быстрым для некоторых массивов, хотя он может плохо масштабироваться до вашей полной проблемы.

Answer 2

Используя np.isin (см. документацию ), вы можете проверить наличие нескольких значений элемента. Например:

import numpy as np 

a = np.array([1,2,3,4])
check_for = np.array([1,2])

locs = np.isin(a, check_for)
# [True,  True, False, False]

np.where(locs)
#[0, 1]

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

В случае, если вам необходимо отслеживать каждое значение элемента отдельно, используйте словарь по умолчанию и перебирайте матрицу.

from collections import defaultdict

tracker = defaultdict(set)

for (row, column), value in np.ndenumerate(arrayElements):
    tracker[value].add(row)

Answer 3

Вы можете попробовать зациклить значения и индексы, используя numpy.ndenumerate и используя Counter, defaultdict или dict, где ключами являются значения в массиве.