Python如何查找唯一条目并从匹配数组中获取最小值

Question

我有一个 numpy 数组，索引:

array([[ 0,  0,  0],
       [ 0,  0,  0],
       [ 2,  0,  2],
       [ 0,  0,  0],
       [ 2,  0,  2],
       [95, 71, 95]])

我有另一个长度相同的数组，称为 distances:

array([  0.98713981,   1.04705992,   1.42340327, 74.0139111 ,
    74.4285216 ,  74.84623217])

indices 中的所有行在 distances 数组中都有一个匹配项。问题是，indices 数组中存在重复项，它们在相应的distances 数组中具有不同的值。我想获得所有三元组索引的最小距离，并丢弃其他的。因此，使用上面的输入，我想要输出:

indicesOUT = 
array([[ 0,  0,  0],
       [ 2,  0,  2],
       [95, 71, 95]])

distancesOUT=
array([  0.98713981,  1.42340327,  74.84623217])

我目前的策略如下:

import numpy as np

indicesOUT = []
distancesOUT = []

for i in range(6):
    for j in range(6):
        for k in range(6):
            if len([s for s in indicesOUT if [i,j,k] == s]) == 0:
                current = np.array([i, j, k])
                ind = np.where((indices == current).all(-1) == True)[0]
                currentDistances = distances[ind]
                dist = np.amin(distances)
                indicesOUT.append([i, j, k])
                distancesOUT.append(dist)

问题是，实际数组每个大约有 400 万个元素，所以这种方法太慢了。最有效的方法是什么？

Answer 1

这本质上是一个分组操作，NumPy 并没有针对它进行优化。幸运的是，Pandas 包有一些非常快速的工具可以适应这个确切的问题。有了上面的数据，我们可以这样做:

import pandas as pd

def drop_duplicates(indices, distances):
    data = pd.Series(distances)
    grouped = data.groupby(list(indices.T)).min().reset_index()
    return grouped.values[:, :3], grouped.values[:, 3]

数据的输出是

array([[  0.,   0.,   0.],
       [  2.,   0.,   2.],
       [ 95.,  71.,  95.]]),
array([  0.98713981,   1.42340327,  74.84623217])

我的基准测试表明，对于 4,000,000 个元素，这应该在大约一秒钟内运行:

indices = np.random.randint(0, 100, size=(4000000, 3))
distances = np.random.random(4000000)
%timeit drop_duplicates(indices, distances)
# 1 loops, best of 3: 1.15 s per loop

如上所述，索引的输入顺序不一定会保留；保持原来的顺序需要多考虑一下。