无循环过滤Numpy数组：按Y列匹配且X列不等筛选并拼接行

问题场景 #

给定Numpy object类型数组（示例如下，实际规模达百万行）：

import numpy as np
x = np.array([['A',1,10],['B',1,20],['C',2,80],['D',3,40],['E',2,50]], dtype=object)

需要实现：

• 遍历每一行，找到数组中第二列（Y列）值相同且第一列（X列）值不同的所有行
• 将这些匹配行直接拼接到当前行之后，最终得到展开后的结果数组

举个例子：第一行Y值为1，匹配到X≠'A'的行['B','1','20']，拼接后这部分结果为['A','1','10'], ['B','1','20']

原代码问题 #

你尝试的代码逻辑混乱，既用了无意义的数组引用y = x，又错误地反转了X列索引做判断，导致返回空数组：

import numpy as np
x = np.array([['A',1,10],['B',1,20],['C',2,80],['D',3,40],['E',2,50]])
y = x
print(x)
print("Result is:",x[np.where(x[:,1] == y[:,1], np.where(x[:,0] != y[:,0][::-1]),False)])

可行解决方案 #

针对百万行规模的数组，优先用Numpy向量化操作提升效率，以下两种方案任选：

方案1：逐行匹配（易理解，效率达标） #

import numpy as np

x = np.array([['A',1,10],['B',1,20],['C',2,80],['D',3,40],['E',2,50]], dtype=object)

# 提取Y列并分组记录索引
y_vals = x[:, 1]
unique_y, indices = np.unique(y_vals, return_inverse=True)

result = []
for i in range(len(x)):
    # 获取当前行同Y值的所有行索引
    same_y_idx = np.where(indices == indices[i])[0]
    # 排除自身，得到匹配行索引
    match_idx = same_y_idx[same_y_idx != i]
    # 先加当前行，再加所有匹配行
    result.append(x[i])
    result.extend(x[match_idx])

# 转换为最终数组
final_array = np.array(result, dtype=object)
print(final_array)

输出结果：

[['A' 1 10]
 ['B' 1 20]
 ['B' 1 20]
 ['A' 1 10]
 ['C' 2 80]
 ['E' 2 50]
 ['D' 3 40]
 ['E' 2 50]
 ['C' 2 80]]

方案2：分组批量处理（更高效，适合超大规模数组） #

先按Y列分组，再批量处理每组内的行，减少重复计算：

import numpy as np

x = np.array([['A',1,10],['B',1,20],['C',2,80],['D',3,40],['E',2,50]], dtype=object)

y_vals = x[:, 1]
unique_y, inverse = np.unique(y_vals, return_inverse=True)

groups = []
for y_idx in range(len(unique_y)):
    # 取出当前Y值对应的所有行
    group_rows = x[inverse == y_idx]
    n_rows = len(group_rows)
    if n_rows == 1:
        # 组内只有一行，直接加入结果
        groups.append(group_rows)
    else:
        # 对组内每一行，拼接组内其他X不同的行
        for row in group_rows:
            others = group_rows[group_rows[:,0] != row[0]]
            groups.append(np.vstack([row, others]))

# 合并所有组的结果
final_array = np.vstack(groups)
print(final_array)

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