使用Numpy Vectorize查找含指定元素子列表索引的报错解决

你之前代码报错有两个核心原因：

1. 回调函数写法错误，numpy.vectorize的lambda表达式漏写了参数x
2. 当所有子列表长度一致时，numpy会自动将输入的嵌套列表转换为二维int数组，numpy.vectorize会遍历数组内的每个int元素，而非遍历外层的子列表对象，因此执行4 in x时x为int类型，触发不可迭代报错。

以下是不同场景下的可用实现方案： #

方案1：原生Python列表推导式（最简洁通用） #

不需要依赖第三方库，代码可读性最高，适合绝大多数场景：

outer_list = [[1,2,3],[1,4,5],[2,7,6,4]]
target = 4
result = [idx for idx, sub_list in enumerate(outer_list) if target in sub_list]
# 输出为 [1, 2]

方案2：提前构建反向索引（高频率查询场景效率最高） #

如果需要对同一个嵌套列表多次查询不同目标元素，提前构建元素到索引的映射，后续查询时间复杂度为O(1)：

from collections import defaultdict

def build_element_index(outer_list):
    index_map = defaultdict(list)
    for idx, sub_list in enumerate(outer_list):
        for elem in sub_list:
            index_map[elem].append(idx)
    return index_map

# 构建索引
outer_list = [[1,2,3],[1,4,5],[2,7,6,4]]
index = build_element_index(outer_list)
# 查询元素4的索引
result = index.get(4, [])
# 输出为 [1, 2]

方案3：修复后的Numpy实现 #

如果必须使用numpy处理，强制指定数组类型为object，避免numpy自动拆分二维数组即可：

import numpy as np

outer_list = [[1,2,3],[1,4,5],[2,7,4]]
target = 4
# 强制创建object类型数组，每个元素为完整子列表
arr = np.array(outer_list, dtype=object)
check_contain = np.vectorize(lambda x: target in x)
result = np.where(check_contain(arr))[0].tolist()
# 输出为 [1, 2]

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