如何改进遍历字典与列表的for循环，获取预期扁平化输出

我写了个for循环，用来遍历my_dictionary的取值和words列表里的元素。需求是：如果words里的单词出现在my_dictionary的某个取值列表中，就把这个取值列表只加一次到结果里，而且最终结果要是扁平化的列表。但运行代码后，结果不仅有重复的嵌套列表，也不是预期的扁平化格式。

原代码 #

my_dictionary = {1: ['apple','dog','cat','bird'],
                 2: ['mouse','rat','elephant','donkey'],
                 3: ['tiger','lion','bear','tortoise']}

words = ['apple', 'dog', 'mouse', 'bear','tortoise']
                
def cen_lst(my_dict, words):
    new_lst = []
    for key, value in my_dict.items():
        for word in words:
            if word in value:
                new_lst.append(value)
    return new_lst  

预期输出 #

['apple','dog','cat','bird', 'mouse','rat','elephant','donkey', 'tiger','lion','bear','tortoise' ]

实际输出 #

[['apple', 'dog', 'cat', 'bird'],
 ['apple', 'dog', 'cat', 'bird'],
 ['mouse', 'rat', 'elephant', 'donkey'],
 ['tiger', 'lion', 'bear', 'tortoise'],
 ['tiger', 'lion', 'bear', 'tortoise']]

问题原因 #

1. 重复添加列表：原代码里，只要words中有一个单词在value列表里，就会把value添加一次。比如['apple','dog',...]这个列表，因为words里有apple和dog两个匹配项，所以被重复添加了两次。
2. 未扁平化列表：直接用append()把整个value列表塞进结果，导致最终是嵌套结构，而不是把列表里的元素逐个展开。

修复后的代码 #

方法一：基础版（清晰易懂） #

my_dictionary = {1: ['apple','dog','cat','bird'],
                 2: ['mouse','rat','elephant','donkey'],
                 3: ['tiger','lion','bear','tortoise']}

words = ['apple', 'dog', 'mouse', 'bear','tortoise']
                
def cen_lst(my_dict, words):
    new_lst = []
    added_values = set()  # 记录已经添加过的列表，避免重复
    for value in my_dict.values():
        # 检查当前列表是否有任意单词匹配words
        has_match = any(word in value for word in words)
        if has_match and value not in added_values:
            new_lst.extend(value)  # 用extend逐个添加元素，实现扁平化
            added_values.add(value)
    return new_lst  

print(cen_lst(my_dictionary, words))

方法二：高效版（适合大数据量） #

把words转成集合，利用集合的交集判断提升效率：

my_dictionary = {1: ['apple','dog','cat','bird'],
                 2: ['mouse','rat','elephant','donkey'],
                 3: ['tiger','lion','bear','tortoise']}

words_set = {'apple', 'dog', 'mouse', 'bear','tortoise'}  # 转成集合加速查询
                
def cen_lst(my_dict, words_set):
    new_lst = []
    for value in my_dict.values():
        # 判断两个集合是否有交集，有交集就添加元素
        if set(value) & words_set:
            new_lst.extend(value)
    return new_lst  

print(cen_lst(my_dictionary, words_set))

修复要点 #

• 用any()判断当前列表是否有匹配项，只要有一个匹配就停止检查该列表的其他单词，避免重复添加。
• 用extend()替代append()，把列表里的元素逐个加入结果，实现扁平化。
• 高效版利用集合的O(1)查询特性，比逐个遍历单词更快，数据量大时优势明显。

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