从含多numpy数组的字典中根据值获取对应键的方法

问题场景 #

现有字典结构如下，值为不等长numpy数组：

import numpy as np
mydict = {'A': np.array([ 1,  5,  6, 10, 13, 17, 20, 22]), 
          'B': np.array([ 3,  8,  9, 15, 16]), 
          'C': np.array([ 0,  2,  4,  7, 11, 12, 14])}

需要根据给定数值（如22）找到对应的字典键，但尝试反转字典时触发TypeError: unhashable type: 'numpy.ndarray'，直接判断数值是否在数组中时触发ValueError: The truth value of an array with more than one element is ambiguous。

错误原因 #

1. 反转字典报错：numpy数组是不可哈希类型，不能作为字典的键，因此无法直接将原字典的键值对反转。
2. 真值模糊错误：直接使用arr == num会返回布尔数组，Python无法直接判断布尔数组的真值，需用.any()确认数组中是否存在匹配值。

解决方案 #

1. 单数值查找 #

遍历字典键值对，用(arr == num).any()判断数值是否在当前数组中：

def get_key_for_num(num):
    for key, arr in mydict.items():
        if (arr == num).any():
            return key
    return None  # 数值不存在时返回None

# 示例调用
print(get_key_for_num(22))  # 输出 'A'
print(get_key_for_num(10))  # 输出 'A'

也可以直接用num in arr（numpy数组原生支持in操作），代码更简洁：

def get_key_for_num(num):
    for key, arr in mydict.items():
        if num in arr:
            return key
    return None

2. 批量数值映射（适合导出数据场景） #

如果需要给大量数据行添加对应键，建议提前生成数值-键的映射字典，避免重复遍历原字典，提升效率：

# 构建数值到键的映射
value_key_map = {}
for key, arr in mydict.items():
    for num in arr:
        value_key_map[num] = key

# 批量处理示例
sample_values = [10, 15, 7, 99]
for val in sample_values:
    print(f"数值{val}对应键：{value_key_map.get(val, '无匹配')}")
# 输出：
# 数值10对应键：A
# 数值15对应键：B
# 数值7对应键：C
# 数值99对应键：无匹配

之后只需通过value_key_map.get(num)即可快速获取对应键，直接添加到数据行中即可导出。

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