如何用Python根据筛选列表过滤距离矩阵获取指定机器间距离

问题背景 #

给定一个10×10的机器间距离矩阵，需要根据动态变化的筛选列表（可能包含子列表）提取指定机器之间的距离子矩阵。例如：

• 原始矩阵：

matrix = [[0, 1, 3, 17, 24, 12, 18, 16, 17, 15],
          [1, 0, 2, 2, 5, 6, 13, 11, 12, 10],
          [3, 2, 0, 1, 6, 12, 18, 12, 17, 15],
          [17, 2, 1, 0, 3, 12, 17, 15, 16, 14],
          [24, 5, 6, 3, 0, 1, 24, 22, 23, 21],
          [12, 6, 12, 12, 1, 0, 12, 10, 11, 9],
          [18, 13, 18, 17, 24, 12, 0, 3, 4, 5],
          [16, 11, 12, 15, 22, 10, 3, 0, 1, 2],
          [17, 12, 17, 16, 23, 11, 4, 1, 0, 1],
          [15, 10, 15, 14, 21, 9, 5, 2, 1, 0]]

• 筛选列表示例：filter_list = [1, 2, 7, 10]（机器编号为1-based）
• 期望输出：

final_matrix = [[0, 1, 18, 15], 
                [1, 0, 13, 10], 
                [18, 13, 0, 5], 
                [15, 10, 5, 0]]

现有代码无法正确处理嵌套列表，且会保留多余元素，需要修正。

问题分析 #

原有代码存在以下问题：

1. 错误假设筛选列表的每个元素都是子列表，当列表为扁平结构时会触发索引错误
2. 仅提取完整行，未筛选对应列，导致结果包含无关元素
3. 未处理筛选列表包含嵌套子列表的情况

解决方案 #

核心步骤：

1. 扁平化筛选列表：将可能嵌套的筛选列表转换为扁平的1-based机器编号列表
2. 转换为0-based索引：将机器编号转换为矩阵对应的索引（矩阵为0-based）
3. 提取子矩阵：遍历筛选后的行索引，提取每行中对应列索引的元素，组成子矩阵

代码实现 #

def flatten_list(lst):
    """扁平化嵌套列表"""
    flattened = []
    for item in lst:
        if isinstance(item, list):
            flattened.extend(flatten_list(item))
        else:
            flattened.append(item)
    return flattened

# 原始矩阵
matrix = [[0, 1, 3, 17, 24, 12, 18, 16, 17, 15],
          [1, 0, 2, 2, 5, 6, 13, 11, 12, 10],
          [3, 2, 0, 1, 6, 12, 18, 12, 17, 15],
          [17, 2, 1, 0, 3, 12, 17, 15, 16, 14],
          [24, 5, 6, 3, 0, 1, 24, 22, 23, 21],
          [12, 6, 12, 12, 1, 0, 12, 10, 11, 9],
          [18, 13, 18, 17, 24, 12, 0, 3, 4, 5],
          [16, 11, 12, 15, 22, 10, 3, 0, 1, 2],
          [17, 12, 17, 16, 23, 11, 4, 1, 0, 1],
          [15, 10, 15, 14, 21, 9, 5, 2, 1, 0]]

# 筛选列表（支持嵌套结构）
filter_list = [1, 2, 7, 10]
# 或者嵌套示例：filter_list = [1, [2, 7], 10]

# 1. 扁平化筛选列表
flat_filter = flatten_list(filter_list)
# 2. 转换为0-based索引
indices = [num - 1 for num in flat_filter]
# 3. 生成子矩阵
final_matrix = [[matrix[i][j] for j in indices] for i in indices]

# 打印结果
print(final_matrix)

测试验证 #

• 当使用filter_list = [1, 2, 7, 10]时，输出与期望一致
• 当使用嵌套列表filter_list = [1, [2, 7], 10]时，输出同样正确

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