Numpy数组按列遍历，基于当前与下一个元素值修改特定元素的实现问题

实现代码 #

import numpy as np

# 原始数组
data = np.array([
     [1,2,3],
     [1,2,3],
     [1,2,101],
     [4,5,111],   
     [4,5,6], 
     [4,5,6], 
     [4,5,101], 
     [4,5,112], 
     [4,5,6], 
     ])

# 提取第三列
col = data[:, 2]
# 构造匹配掩码：当前元素为6，下一个相邻元素为101
mask = (col[:-1] == 6) & (col[1:] == 101)
# 替换匹配位置的元素
col[np.where(mask)[0]] = 10001

# 输出处理后的数组
print(data)

运行结果 #

[[    1     2     3]
 [    1     2     3]
 [    1     2   101]
 [    4     5   111]
 [    4     5     6]
 [    4     5 10001]
 [    4     5   101]
 [    4     5   112]
 [    4     5     6]]

原理解释 #

• 提取数组第三列单独处理，避免操作整个数组浪费性能
• 用切片col[:-1]取到除最后一个元素外的所有第三列元素（所有存在下一个相邻元素的位置），col[1:]取到除第一个元素外的所有第三列元素（对应所有前一个位置的下一个元素）
• 两个条件做按位与得到匹配掩码，掩码为True的位置就是需要替换的位置索引
• 直接通过索引批量替换匹配位置的元素为10001，全程采用numpy向量化操作，无需循环，性能远高于遍历数组的实现

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