Python实现两个嵌套列表子列表首元素匹配并返回首列表对应项

核心思路 #

• 先提取list2中所有子列表的首元素存入集合，集合的成员查询时间复杂度为O(1)，相比遍历列表查询效率更高，适合数据量较大的场景
• 遍历list1的子列表，仅保留首元素存在于上述集合中的条目，即可得到预期结果

代码实现 #

def function(list1, list2):
    # 提取list2所有非空子列表的首元素存入集合，避免空列表索引报错
    list2_first_keys = {sublist[0] for sublist in list2 if sublist}
    # 筛选list1中首元素匹配的非空子列表
    return [sublist for sublist in list1 if sublist and sublist[0] in list2_first_keys]

效果验证 #

使用示例提供的测试数据运行：

list1 = [['a','b','c'], ['e','f','g'], ['i','j','k']]
list2 = [['a','m','n'], ['i','m','k'], ['q','r','s']]
print(function(list1, list2))

运行输出：

[['a', 'b', 'c'], ['i', 'j', 'k']]

与预期结果完全一致。

补充说明 #

如果可以确定输入的两个嵌套列表中所有子列表都不为空，可以把代码中的if sublist判断去掉，简化写法。

内容的提问来源于stack exchange，提问作者Soumya Ranjan Sahoo