Python面试题求助：多列表字符串连接计数的高性能方案优化

我刚参加了一家创业公司的Python技术面试，15分钟内给出了解决方案，但推测面试官更关注性能优化，希望找到效率更高的实现方式。

题目描述 #

给定一个字符串的列表的列表，例如：

lst = [['id_2','id_3'],['id_3','id_4','id_2'],['id_5']]

需要返回每个字符串的连接数：同一个子列表内的所有字符串互为连接（连接关系去重），最终返回每个字符串对应的唯一连接对象数量。示例预期结果为：
{'id_2':2, 'id_3':2 , 'id_4':2 , 'id_5':0}

我的初始解决方案 #

lst = [['id_2','id_3'],['id_3','id_4','id_2'],['id_5']]

# Answer:- {id_2:2, id_3:2 , id_4:2 , id_5:0}
class sln:
    def __init__(self, lst: list):
        self._data = {} # key: id_X, value: list of connections [id_y, id_z]
        self._build_data(lst)

    def count_relations(self):
        result = {}
        for entry in self._data:
            result[entry] = len(self._data[entry])
        return result

    def _build_data(self, lst: list):
        for outter_lst in lst:
            for item in outter_lst:
                if item not in self._data.keys():
                    self._data[item] = []
                self._append_relevant_keys(item, outter_lst)

    def _append_relevant_keys(self, item: str, outter_lst: list):
        for entry in outter_lst:
            if entry not in self._data[item] and item != entry:
                self._data[item].append(entry)


if __name__ == '__main__':
    print(sln(lst).count_relations())

初始方案的核心问题在于：用列表存储连接关系时，每次判断entry not in self._data[item]是O(n)的线性遍历，当数据量较大时，重复的判断会导致性能急剧下降。面试时我提到可以将列表替换为集合来提升查询速度，下面是具体的优化方案。

性能优化方案 #

核心优化思路 #

1. 用集合替代列表存储连接关系：集合的成员查询、插入操作都是O(1)时间复杂度，避免了原方案中线性遍历判断的开销，同时自动去重，无需手动处理重复连接。
2. 批量处理子列表：将子列表提前转为集合，通过集合差集直接获取当前元素的所有连接对象，减少嵌套循环的冗余操作。
3. 利用collections.defaultdict简化初始化：自动为新的ID创建空集合，省去手动判断和初始化的步骤。

优化后的代码 #

from collections import defaultdict

lst = [['id_2','id_3'],['id_3','id_4','id_2'],['id_5']]

class OptimizedSln:
    def __init__(self, lst: list):
        self._connections = defaultdict(set)  # 自动初始化空集合
        self._build_connections(lst)

    def count_relations(self):
        # 用字典推导式直接生成结果，简洁高效
        return {id: len(connections) for id, connections in self._connections.items()}

    def _build_connections(self, lst: list):
        for group in lst:
            group_set = set(group)
            for item in group_set:
                # 批量添加当前组内除自身外的所有连接对象
                self._connections[item].update(group_set - {item})

if __name__ == '__main__':
    print(OptimizedSln(lst).count_relations())

性能对比 #

• 原方案时间复杂度：O(M*N²)，其中M是子列表数量，N是子列表平均长度。每个元素需要遍历子列表N次，每次判断是否在列表中是O(N)操作。
• 优化后方案时间复杂度：O(M*N)，集合的转置、差集、更新操作均为线性时间，整体开销与数据总量成正比，性能提升显著，尤其适合处理大规模数据。

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