获取许多掩码的每个掩码 np.array 的最大值的最快方法？

Question

我有两个相同形状的 numpy 数组。一个包含我感兴趣的信息，另一个包含一堆可以用作掩码值的整数。

本质上，我想遍历每个唯一整数以获取数组的每个掩码，然后使用此掩码过滤主数组并找到过滤后数组的最大值。

为简单起见，假设数组是：

arr1 = np.random.rand(10000,10000)
arr2 = np.random.randint(low=0, high=1000, size=(10000,10000))

现在我正在这样做：

maxes = {}
ids = np.unique(arr2)
for id in ids:
    max_val = arr1[np.equal(arr2, id)].max()
    maxes[id] = max_val

我的阵列要大得多，而且速度非常慢，我正在努力寻找一种更快的方法来做到这一点......也许有某种我不知道的创造性方法，非常感谢任何帮助。

编辑

假设大多数 arr2 实际上是 0，我不关心 0 id，是否可以通过从搜索中删除整个块来加快速度？？

即

arr2[:, 0:4000] = 0

只返回 ids > 0 的最大值？？

非常感谢..

Answer 1

通用的基于 bin 的缩减策略

下面列出了一些方法来解决我们需要执行基于 bin 的归约操作的场景。所以，基本上我们得到了两个数组，我们需要使用一个作为 bin，另一个作为值，并减少第二个。

方法#1：一种策略是根据arr2 对arr1 进行排序。一旦我们以相同的顺序对它们进行排序，我们就会找到组开始和停止索引，然后使用适当的ufunc.reduceat，我们执行基于切片的归约操作。就是这样！

这是实现 -

def binmax(bins, values, reduceat_func):
    ''' Get binned statistic from two 1D arrays '''
    sidx = bins.argsort()
    bins_sorted = bins[sidx]
    grpidx = np.flatnonzero(np.r_[True,bins_sorted[:-1]!=bins_sorted[1:]])
    max_per_group = reduceat_func(values[sidx],grpidx)
    out = dict(zip(bins_sorted[grpidx], max_per_group))
    return out

out = binmax(arr2.ravel(), arr1.ravel(), reduceat_func=np.maximum.reduceat)

它适用于具有相应 ufunc.reduceat 方法的 ufunc。

方法#2：我们还可以利用scipy.stats.binned_statistic，它基本上是一个通用实用程序来执行一些基于分箱数组值的常见归约操作 -

from scipy.stats import binned_statistic

def binmax_v2(bins, values, statistic):
    ''' Get binned statistic from two 1D arrays '''
    num_labels = bins.max()+1
    R = np.arange(num_labels+1)
    Mx = binned_statistic(bins, values, statistic=statistic, bins=R)[0]
    idx = np.flatnonzero(~np.isnan(Mx))
    out  = dict(zip(idx, Mx[idx].astype(int)))
    return out

out = binmax_v2(arr2.ravel(), arr1.ravel(), statistic='max')