Python numpy：为3个其他列的每个唯一元组有效地获取包含min值的行

Posted 2021-05-06

我有一些数据存储在列表列表中（约200,000行×6列的列表）。

我需要得到以下数据子集：对于列[1,2,4]中的每个唯一值集，我需要找到列0的最小值并且仅保留该行。

我必须在旧的numpy 1.10中执行此操作（不要问...），因此np.unique（）中没有'axis = 0'选项。

以下示例运行并生成正确的输出，但速度非常慢。这看起来很基本，所以我觉得（缺乏）速度一定是我的错。

# S-L-O-W way to get the desired output:
import numpy as np

# Example dataset
data = [[1, 1, 1, 'a', 1],
        [0, 1, 1, 'b', 1],
        [0, 3, 1, 'c', 4],
        [3, 1, 1, 'd', 1],
        [4, 3, 1, 'e', 4]]

desired_output = [[0, 1, 1, 'b', 1],
                  [0, 3, 1, 'c', 4]]

# Currently coding on a geriatric machine with numpy pre-version 1.13 and no ability to upgrade,
# so np.unique() won't take an axis argument. The next few hack lines of code get around this with strings...
tuples_str = []
tuples_raw = [[datarow[jj] for jj in [1,2,4]]  for datarow in data ]
for datarow in data:
    one_tuple = [datarow[jj] for jj in [1,2,4]]
    tuples_str.append( '_'.join([str(ww) for ww in one_tuple]) )

# Numpy unique on this data subset with just columns [1,2,4] of original data
unq, unq_inv, unq_cnt = np.unique(tuples_str, return_inverse=True, return_counts=True)

# Storage
output = []

# Here's the painfully slow part:
# Iterate over each subset of data where rows take the value in one unique tuple (i.e. columns [1,2,4] are identical)
for ii, idx in enumerate(np.unique(unq_inv)):

    # Get the rows that have the same values in columns [1,2,4]
    all_matches_thistuple = [row for ii, row in enumerate(data) if unq_inv[ii]==idx]

    # Find the index of the row with the minimum value for column 0
    first_line_min_idx = np.argmin([int(row1[0]) for row1 in all_matches_thistuple])

    # Save only that row
    output.append(all_matches_thistuple[first_line_min_idx])
print(output)

答案

如果您从列表列表开始，您可以使用普通的Python轻松完成此任务，并且它将具有高性能。实际上，你正在使用带有numpy dtype的object，所以我怀疑你使用内置例程获得很少的性能，因为你丢失了数据局部性（并且基本上留下了相当于糟糕的Python list对象）。相反，你可以在线性时间内完成这个（不计算数据的初始排序，它将是O（n * logN），但它将使用Python的timsort，所以它实际上会非常快），只做一对传递数据：

In [1]: data = [[1, 1, 1, 'a', 1],
   ...:         [0, 1, 1, 'b', 1],
   ...:         [0, 3, 1, 'c', 4],
   ...:         [3, 1, 1, 'd', 1],
   ...:         [4, 3, 1, 'e', 4]]
   ...:

In [2]: from operator import itemgetter

In [3]: group_key = itemgetter(1,2,4)

In [4]: data.sort(key=group_key)

那简单地说：

In [6]: first = itemgetter(0)

In [7]: result = []

In [8]: from itertools import groupby
   ...: for _, g in groupby(data, group_key):
   ...:     result.append(min(g, key=first))
   ...:

In [9]: result
Out[9]: [[0, 1, 1, 'b', 1], [0, 3, 1, 'c', 4]]

另一种方法是使用defaultdict构建辅助数据结构。这是对非排序数据进行分组的惯用方法。如果您希望能够将这些值分组，这可能很有用：

In [10]: from collections import defaultdict

In [11]: grouper = defaultdict(list)

In [12]: data = [[1, 1, 1, 'a', 1],
    ...:         [0, 1, 1, 'b', 1],
    ...:         [0, 3, 1, 'c', 4],
    ...:         [3, 1, 1, 'd', 1],
    ...:         [4, 3, 1, 'e', 4]]

In [13]: for row in data:
    ...:     _,x,y,_, z = row
    ...:     grouper[(x,y,z)].append(row)
    ...:

In [14]: grouper
Out[14]:
defaultdict(list,
            {(1, 1, 1): [[1, 1, 1, 'a', 1],
              [0, 1, 1, 'b', 1],
              [3, 1, 1, 'd', 1]],
             (3, 1, 4): [[0, 3, 1, 'c', 4], [4, 3, 1, 'e', 4]]})

In [15]: first = itemgetter(0)

In [16]: [min(group, key=first) for group in grouper.values()]
Out[16]: [[0, 1, 1, 'b', 1], [0, 3, 1, 'c', 4]]

另一答案

这是Pandas的一种方式，如果你可以让它工作：

df = pd.DataFrame(data).sort_values(0).drop_duplicates([1, 2, 4]).values

结果

[[0 1 1 'b' 1]
 [0 3 1 'c' 4]]

说明

您的问题可以简化为：

• 按列0排序，ascending=True是默认值。
• 按列[1,2,4]删除重复的行。
• pd.DataFrame.values提取底层的numpy数组。