如何高效提取二维numpy数组每行前5个相邻不重复元素

实现方案 #

以下是完全基于numpy向量化操作的无循环实现，20000行规模的数组可以毫秒级完成运算：

import numpy as np

# 示例输入数组
a = np.array([[193,  64,  64, 139, 180, 180, 104, 152,  69,  22, 192,  92],
              [  1,  36, 156, 152,  152,  37,  46, 143, 141, 114,  25, 134],
              [110,  96,  52,  53,  35, 147,   3, 116,  20,  11, 137,   5]])

# 1. 构造相邻元素不等的掩码，第一列默认保留
mask = np.ones(a.shape, dtype=bool)
mask[:, 1:] = a[:, 1:] != a[:, :-1]

# 2. 按行计算掩码累积和，标记截止到当前位置的相邻不重复元素个数
cumsum_mask = np.cumsum(mask, axis=1)

# 3. 筛选前5个不重复元素并整理为二维数组
result = a[cumsum_mask <= 5].reshape(-1, 5)

输出验证 #

打印result即可得到预期结果：

array([[193,  64, 139, 180, 104],
       [  1,  36, 156, 152,  37],
       [110,  96,  52,  53,  35]])

原理解释 #

• 掩码mask会标记出所有相邻去重后需要保留的元素位置，第一列没有左侧元素默认全保留
• 按行累积求和的cumsum_mask每个位置的值，代表截止到当前位置本行已经出现的相邻不重复元素总数
• 筛选累积和≤5的元素后按行重排，即可直接得到每行前5个相邻不重复元素的结果

内容的提问来源于stack exchange，提问作者Ernest S Kirubakaran