基于自定义规则对Python嵌套列表排序并覆盖所有元素

问题说明 #

原始嵌套列表：

[[1, 0], [2, 1], [5, 4], [1, 3], [4, 1], [3, 2], [0, "<NA>"]]

排序规则：

• 第二个元素为<NA>的子列表必须排在首位
• 后续子列表需满足：前一个子列表的第一个元素与后一个子列表的第二个元素相等，例如[0, "<NA>"], [1, 0], [2, 1]
• 排序后的列表必须包含原列表所有元素

预期输出：

[[0, "<NA>"], [1, 0], [2, 1], [3, 2], [1, 3], [4, 1], [5, 4]]

注意：若在[1, 0]之后选择[4, 1]，会无法覆盖所有元素，因此该选择不符合要求。

实现方案 #

我们可以通过构建映射+回溯查找的方式实现，确保找到能覆盖所有元素的有效路径：

def sort_nested_list(lst):
    # 构建元素映射：以子列表的第二个元素为键，对应子列表为值
    from collections import defaultdict
    mapping = defaultdict(list)
    start = None
    
    for item in lst:
        if item[1] == "<NA>":
            start = item
        else:
            mapping[item[1]].append(item)
    
    # 回溯寻找完整路径
    def backtrack(current_path, remaining_mapping):
        # 若所有元素已加入路径，返回结果
        if sum(len(v) for v in remaining_mapping.values()) == 0:
            return current_path
        
        last_first_val = current_path[-1][0]
        # 没有可匹配的下一个元素，回溯
        if last_first_val not in remaining_mapping or not remaining_mapping[last_first_val]:
            return None
        
        # 尝试所有可能的下一个元素
        for idx, next_item in enumerate(remaining_mapping[last_first_val]):
            # 复制映射避免修改原数据
            new_mapping = defaultdict(list)
            for k, v in remaining_mapping.items():
                new_mapping[k] = v.copy()
            # 移除当前选中的元素
            new_mapping[last_first_val].pop(idx)
            if not new_mapping[last_first_val]:
                del new_mapping[last_first_val]
            
            result = backtrack(current_path + [next_item], new_mapping)
            if result:
                return result
        return None
    
    return backtrack([start], mapping)

# 测试示例
original_list = [[1, 0], [2, 1], [5, 4], [1, 3], [4, 1], [3, 2], [0, "<NA>"]]
print(sort_nested_list(original_list))

代码说明 #

1. 先定位起始元素（第二个值为<NA>的子列表），同时构建映射表，快速查找可匹配的后续元素
2. 回溯函数会尝试每一种可能的路径分支，直到找到包含所有元素的有效路径
3. 每次尝试分支时都会复制映射表，避免影响其他分支的查找逻辑

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