python - 是否有比 np.where 更快的替代方案来确定索引？

Question

我有一个像这样的数组:

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

例如，我需要知道 arrayElements 等于 1 的索引。

现在，我正在做:

行、列 = np.where(arrayElements == 1)

这可行，但我是在一个循环中执行此操作，循环遍历所有可能的元素值，在我的例子中，它是 1-500,000+。运行需要 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]})

对于某些数组来说，第二种方法可能更快，但它可能无法很好地解决您的整个问题。