使用循环查找两个numpy数组匹配值对应索引的代码修正问题

原代码问题分析 #

• 双重循环逻辑冗余，且没有收集符合条件的data数组索引，只是打印了单元素比对的np.where返回值，无法得到预期结果
• 定义了空列表a但未实际使用存储匹配索引
• 针对单值调用np.where得到的[0]无实际意义，不是你需要的data数组的位置索引

最优实现（NumPy向量化方案，无需循环） #

直接用np.isin批量判断data中元素是否存在于data1，再用np.where提取匹配位置的索引即可：

import numpy as np

data = ['2007-01-03', '2007-01-04', '2007-01-05','2007-01-03', '2007-01-04', '2007-01-05']
data1 = ['2007-01-04', '2007-01-05']

data = np.array(data)
data1 = np.array(data1)

# 批量判断匹配位置，转换为列表输出
match_index = np.where(np.isin(data, data1))[0].tolist()
print(match_index)

运行输出：
[1, 2, 4, 5]

原循环逻辑修正版 #

如果你需要保留循环写法，去掉多余的内层循环，直接判断当前元素是否在data1中，收集对应索引即可：

import numpy as np

a = []
data = ['2007-01-03', '2007-01-04', '2007-01-05','2007-01-03', '2007-01-04', '2007-01-05']
data1 = ['2007-01-04', '2007-01-05']

data = np.array(data)
data1 = np.array(data1)

for idx, day in enumerate(data):
    if day in data1:
        a.append(idx)
print(a)

运行输出和上述方案一致。

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