如何高效过滤二维NumPy数组中含前置行元素的重复行？

问题描述 #

需要过滤二维NumPy数组的行：若某一行中任意元素已在之前保留的行中出现过，则剔除该行。示例如下：

import numpy as np

# 输入数组
arr = np.array([
    [4580, 4581, 4657, 4658],
    [4580, 4581, 4657, 4659],  # 含已出现的4580，剔除
    [4650, 4652, 4654, 4655],
    [4651, 4655, 4652, 4656]   # 含已出现的4652，剔除
])

# 期望输出
# array([[4580, 4581, 4657, 4658],
#        [4650, 4652, 4654, 4655]])

原代码在小数据下有效，但大数据时因重复计算唯一值导致效率极低：

check = np.array([not(np.in1d(a, np.unique(arr[:i])).any()) for i,a in enumerate(arr)])
arr[check]

高效解决方案 #

核心思路是用集合跟踪已出现元素，避免重复计算历史元素的唯一值：

1. 初始化一个空集合，用于记录所有已保留行的元素
2. 遍历数组的每一行，检查该行是否有元素在集合中
3. 若该行无元素在集合中，则保留该行，并将该行所有元素加入集合
4. 最后将保留的行转为NumPy数组

代码实现 #

import numpy as np

def filter_rows_by_element_history(arr):
    seen = set()
    keep_rows = []
    for row in arr:
        # 检查当前行是否有元素已被记录
        if not any(x in seen for x in row):
            keep_rows.append(row)
            # 将当前行所有元素加入已记录集合
            seen.update(row)
    return np.array(keep_rows)

# 测试
arr = np.array([
    [4580, 4581, 4657, 4658],
    [4580, 4581, 4657, 4659],
    [4650, 4652, 4654, 4655],
    [4651, 4655, 4652, 4656]
])

result = filter_rows_by_element_history(arr)
print(result)

效率说明 #

• 原方法的时间复杂度为O(n²m)（每次循环都要对前i行去重，去重操作是O(im)），大数据下会急剧变慢
• 新方法的时间复杂度为O(n*m)，集合的查询和更新都是O(1)平均时间复杂度，处理百万级行数也能保持高效

补充优化（纯NumPy实现，适合极端大数据） #

如果需要完全用NumPy操作避免Python循环（当每行元素数极大时），可以用以下方法：

import numpy as np

def filter_rows_numpy(arr):
    seen = np.array([], dtype=arr.dtype)
    mask = np.ones(arr.shape[0], dtype=bool)
    for i, row in enumerate(arr):
        # 检查当前行元素是否在已记录数组中
        overlap = np.isin(row, seen)
        if overlap.any():
            mask[i] = False
        else:
            seen = np.concatenate([seen, row])
    return arr[mask]

注：此方法在元素数量极大时，np.isin的性能略低于Python集合，但纯NumPy操作适合某些需要避免GIL的场景。

内容的提问来源于stack exchange，提问作者Kristof