更快的替代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 gig的内存.我假设这不是一个选择.

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

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]

这利用稀疏矩阵构造函数中的快速代码以有用的方式组织数据,构造稀疏矩阵,其中行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最终是一种非常紧凑和有效的方式来存储所有相关信息供以后使用,如问题的附加部分所述.希望这很有用!

编辑:我意识到初始版本中存在一个错误:如果数据中没有出现任何值,则失败:现在已经修复了.

Answer 2

我正在考虑这一点,并意识到使用Pandas解决这个问题的方法更直观(但速度稍慢)groupby().考虑一下:

import numpy as np
import pandas as pd

def get_indices_pandas(data):
    d = data.ravel()
    f = lambda x: np.unravel_index(x.index, data.shape)
    return pd.Series(d).groupby(d).apply(f)

这将返回与get_indices_simple我之前的答案相同的结果:

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_pandas(data_small)))
# True

而且这种Pandas方法比不那么直观的矩阵方法略慢:

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

%time ind = get_indices_simple(data)
# CPU times: user 14.2 s, sys: 665 ms, total: 14.8 s
# Wall time: 14.9 s

%time ind = get_indices_sparse(data)
# CPU times: user 842 ms, sys: 277 ms, total: 1.12 s
# Wall time: 1.12 s

%time ind = get_indices_pandas(data)
# CPU times: user 1.16 s, sys: 326 ms, total: 1.49 s
# Wall time: 1.49 s