如何高效按给定列表值过滤二维np array且保留原元素位置

最优实现方案 #

直接使用Numpy内置向量化方法实现，全程无Python级循环，大数组场景下性能远超列表推导式实现，同时天然保留元素原有位置顺序：

import numpy as np

# 示例数据
b = np.array([['a', 'b', 'c', 'd'], ['b', 'a', 'c', 'd'], ['c', 'b', 'a', 'd'], ['a', 'd', 'c', 'b']])
values_to_stay_in_b = ['a', 'b']

# 生成布尔掩码，标记所有需要保留的元素
mask = np.isin(b, values_to_stay_in_b)
# 过滤后按原数组行数调整维度，得到最终结果
result = b[mask].reshape(b.shape[0], -1)

运行后得到的result和你期望的输出完全一致：

array([['a', 'b'],
       ['b', 'a'],
       ['b', 'a'],
       ['a', 'b']], dtype='<U1')

特殊场景兼容 #

如果你的实际场景中每行符合保留条件的元素数量不固定，无法直接reshape为规整二维数组，可以用以下方式得到每行的过滤结果列表：

row_keep_count = mask.sum(axis=1)
result = np.split(b[mask], np.cumsum(row_keep_count))[:-1]

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