如何高效查找Python二维列表中元素的子列表索引？

问题描述 #

需要处理二维列表的元素索引映射：list2是list1元素重排后的二维列表，需生成与两者结构一致的list3，其中每个元素为list1对应元素在list2中所属子列表的索引。现有四层for循环及基于np.argwhere的实现效率极低，无法处理十万级规模的列表，寻求更高效的实现方案。

示例 #

list1 = [[0,1],[2,3],[4,5],[6,7,8],[9,10,11],[12,13,14]]
list2 = [[8,2],[1,9],[6,4],[13,5,0],[14,7,10],[12,11,3]]

期望输出：

list3 = [[3,1],[0,5],[2,3],[2,4,0],[1,4,5],[5,3,4]]

现有低效实现 #

1st 实现（四层循环） #

for index1, sublist1 in enumerate(list1):
    for i in range(len(sublist1)):
        for index2, sublist2 in enumerate(list2):
            for j in range(len(sublist2)):
                if sublist2[j] == sublist[1]:  # 注：原代码存在笔误，应为sublist1[i]
                    list1[index1][i] = index1
                    list2[index2][j] = index2

2nd 实现（基于np.argwhere） #

for index1, sublist1 in enumerate(list1):
    for i in range(len(sublist1)):
        where = np.argwhere(list2 == sublist1[i])
        list1[index1][i] = where[0, 0]

高效实现方案 #

核心思路是先构建元素到list2子列表索引的映射字典，利用字典O(1)的查找效率，将整体时间复杂度降至线性级别，完全适配十万级规模数据。

代码实现 #

list1 = [[0,1],[2,3],[4,5],[6,7,8],[9,10,11],[12,13,14]]
list2 = [[8,2],[1,9],[6,4],[13,5,0],[14,7,10],[12,11,3]]

# 第一步：构建元素与list2子列表索引的映射字典
element_to_group = {}
for group_idx, sublist in enumerate(list2):
    for num in sublist:
        element_to_group[num] = group_idx

# 第二步：遍历list1生成list3
list3 = [[element_to_group[num] for num in sublist] for sublist in list1]

print(list3)
# 输出结果：[[3, 1], [0, 5], [2, 3], [2, 4, 0], [1, 4, 5], [5, 3, 4]]

效率说明 #

• 字典构建阶段：遍历list2所有元素，时间复杂度为O(M)（M为list2总元素数）
• 生成list3阶段：遍历list1所有元素，时间复杂度为O(N)（N为list1总元素数）
• 整体时间复杂度为O(M+N)，相比四层循环的O(NM)和np.argwhere的O(NM)，效率提升几个数量级，完全能处理十万级甚至更大规模的数据。

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