numpy.where的更快替代品？

Question

我有一个3d数组，其中填充了从0到N的整数。我需要一个与该数组等于1、2、3，… N的位置对应的索引列表。我可以使用np.where进行如下操作：

N = 300
shape = (1000,1000,10)
data = np.random.randint(0,N+1,shape)
indx = [np.where(data == i_id) for i_id in range(1,data.max()+1)]

但这很慢。根据这个问题 快速python
numpy在哪里功能？

应该可以大大加快索引搜索的速度，但是我无法将那里提出的方法转移到我获取实际索引的问题上。加快上述代码的最佳方法是什么？

作为附加组件：我想稍后存储索引，使用np.ravel_multi_index可以将大小从保存3个索引减少为仅1个，即使用：

indx = [np.ravel_multi_index(np.where(data == i_id), data.shape) for i_id in range(1, data.max()+1)]

这更接近于Matlab的find函数。可以直接将其合并到不使用np.where的解决方案中吗？

Answer 1

我认为解决此问题的标准向量化方法最终将占用大量内存，对于int64数据，将需要O（8 * N * data.size）字节，对于上面给出的示例，则需要约22
gigs的内存。我假设这不是一个选择。

您可以通过使用稀疏矩阵存储唯一值的位置来取得一些进展。例如：

import numpy as np
from scipy.sparse import csr_matrix

def compute_M(data):
    cols = np.arange(data.size)
    return csr_matrix((cols, (data.ravel(), cols)),
                      shape=(data.max() + 1, data.size))

def get_indices_sparse(data):
    M = compute_M(data)
    return [np.unravel_index(row.data, data.shape) for row in M]

这利用了稀疏矩阵构造函数中的快速代码来以一种有用的方式来组织数据，从而构造了一个稀疏矩阵，其中rowi仅包含扁平化数据等于的索引i。

为了进行测试，我还将定义一个执行简单方法的函数：

def get_indices_simple(data):
    return [np.where(data == i) for i in range(0, data.max() + 1)]

对于相同的输入，这两个函数给出相同的结果：

data_small = np.random.randint(0, 100, size=(100, 100, 10))
all(np.allclose(i1, i2)
    for i1, i2 in zip(get_indices_simple(data_small),
                      get_indices_sparse(data_small)))
# True

稀疏方法比数据集的简单方法快一个数量级：

data = np.random.randint(0, 301, size=(1000, 1000, 10))

%time ind = get_indices_simple(data)
# CPU times: user 14.1 s, sys: 638 ms, total: 14.7 s
# Wall time: 14.8 s

%time ind = get_indices_sparse(data)
# CPU times: user 881 ms, sys: 301 ms, total: 1.18 s
# Wall time: 1.18 s

%time M = compute_M(data)
# CPU times: user 216 ms, sys: 148 ms, total: 365 ms
# Wall time: 363 ms

稀疏方法的另一个好处是，矩阵M最终成为存储所有相关信息以供以后使用的非常紧凑而有效的方式，如问题的附加部分所述。希望这是有用的！

编辑：我意识到初始版本中存在一个错误：如果范围中的任何值未出现在数据中，它将失败：现在已在上面修复。