Numpy for 循环向量化：使用 np.all 广播和测试独特元素

Question

工作循环，预期结果

我正在尝试对具有非常大数据集的代码中的慢速 for 循环进行矢量化，以根据测试删除重复项。结果应该只保留前 3 个元素是唯一的元素，并且第 4 个元素是所有重复项中最大的元素。例如

in = np.array(((0, 12, 13, 1), (0, 12, 13, 10), (1, 12, 13, 2)))

应该变成

out = np.array(((0, 12, 13, 10), (1, 12, 13, 2)))

使用 for 循环实现这一点很简单，但正如我所提到的，它非常慢。

unique = np.unique(in[:, :3], axis=0)
out = np.empty((0, 4))
for i in unique:
    out = np.vstack((out, np.hstack((i[:], np.max(in[np.all(in[:, :3] == i[:], axis=1)][:, 3])))))

我尝试过的 (1)

当我尝试通过将每个 i[:] 替换为 unique[np.arange(unique.shape[0])] 来删除带有索引的 for 循环时：

out = np.vstack((out, np.hstack((unique[np.arange(unique.shape[0])], np.max(in[np.all(in[:, :3].astype(int) == unique[np.arange(unique.shape[0])], axis=1)][:, 3])))))

Numpy 抱怨输入形状与所有：

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<__array_function__ internals>", line 6, in all
  File "/usr/local/lib/python3.6/dist-packages/numpy/core/fromnumeric.py", line 2351, in all
    return _wrapreduction(a, np.logical_and, 'all', axis, None, out, keepdims=keepdims)
  File "/usr/local/lib/python3.6/dist-packages/numpy/core/fromnumeric.py", line 90, in _wrapreduction
    return ufunc.reduce(obj, axis, dtype, out, **passkwargs)
numpy.AxisError: axis 1 is out of bounds for array of dimension 0

我尝试过的 (2)

根据 StackOverflow 在输入此问题时的建议 (Broadcasting/Vectorizing inner and outer for loops in python/NumPy)：

newout = np.vstack((newout, np.hstack((tempunique[:, None], np.max(inout[np.all(inout[:, :3].astype(int) == tempunique[:, None], axis=1)][:, 3])))))

我收到一个错误，抱怨输入和输出之间的大小不匹配：

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
IndexError: boolean index did not match indexed array along dimension 0; dimension is 3 but corresponding boolean dimension is 2

重述问题

是否有正确的方法来广播我的索引以消除 for 循环？

Answer 1

我对您的用例了解不多，无法确定是否值得引入 Pandas，但在 Pandas 中有效地这样做只需要几行代码：

import numpy as np
import pandas as pd

in_array = np.array(((0, 12, 13, 1), (0, 12, 13, 10), (1, 12, 13, 2)))
in_df = pd.DataFrame(in_array)


# group by unique combinations of the 0th, 1st, and 2nd columns, then take the
# max of the 3rd column in each group. `reset_index` change cols 0-2 from index
# back to normal columns
out_df = in_df.groupby([0, 1, 2])[3].max().reset_index()
out_array = out_df.values

print(out_array)
# Output:
# [[ 0 12 13 10]
#  [ 1 12 13  2]]

一个简单的时序测试表明，处理一个 100000 行随机生成的输入数组，使用 Pandas 需要 0.0117 秒，使用 for 循环实现需要 2.6103 秒。