Python中查找二维数组b的元素在排序后一维数组a中的索引位置

原有代码的错误点 #

• 函数定义语法错误：def findall_index(b, a ) 末尾缺少冒号
• 循环逻辑颠倒：你需要遍历二维数组b的每个元素，而非遍历一维数组a
• 判断逻辑错误：原代码判断条件b[i][j] == a是用b的单个元素和整个a数组比较，永远不会成立
• 排序后未赋值：你执行了np.sort(a)但没有把结果回写给a，导致后续使用的a还是未排序状态
• 函数调用参数错误：调用时传入的0和未定义的result变量，完全不符合函数要求的入参格式
• 类型不统一问题：b中同时存在数字0和字符串类型的邮编，判断a != '0'时会因为类型不一致导致过滤失效

修改后实现代码 #

先构建元素到索引的映射字典，大幅提升查找效率，避免每次查找都遍历整个a数组：

import pandas as pd
import numpy as np

# 原有数据读取逻辑保持不变
cols = [2,4,6,8,10,12,14,16,18]
df = pd.read_csv('mywork.csv')
df1 = df.iloc[:, cols]
b = np.array(df1, dtype=str) # 统一转成字符串类型，避免类型不一致导致判断错误

# 处理得到排序后的a数组
a = np.concatenate(b)
a = a[(a != 'nan') & (a != '0')] # 一次性剔除无效值，写法更简洁
a = np.sort(a) # 排序后回写给a，修正原代码漏赋值的问题

# 构建元素到索引的映射字典，查找效率更高
value_to_idx = {val: idx for idx, val in enumerate(a)}

# 实现查找函数
def findall_index(b, value_map):
    result = []
    # 遍历b的行和列
    for i in range(b.shape[0]):
        for j in range(b.shape[1]):
            current_val = b[i][j]
            # 只保留在a中存在的元素的位置，无效值可以选择跳过或者标记为-1
            if current_val in value_map:
                # 结果格式可自定义，这里返回(原二维数组行号, 原二维数组列号, 排序后a中的索引)
                result.append((i, j, value_map[current_val]))
    return result

# 调用函数输出结果
res = findall_index(b, value_to_idx)
print(res)

结果说明 #

返回的结果中每个元素是三元组(b的行索引, b的列索引, 对应元素在排序后a中的索引)，如果你只需要a中的索引，或者需要保留无效值的位置标记，调整append的内容即可。

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