numpy.where的更快替代品？

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

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

import numpy as npfrom scipy.sparse import csr_matrixdef 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]

这利用了稀疏矩阵构造函数中的快速代码来以一种有用的方式来组织数据，从而构造了一个稀疏矩阵，其中row

i

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

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