如何从主矩阵中按指定行索引提取两个子矩阵？

需求说明 #

给定主矩阵w，输入行/列索引数组a（示例中为[0,1]），需要：

• subMatrixA：选取a中指定的行，且仅保留这些行里列索引属于a的元素
• subMatrixB：选取a中指定的行，保留这些行里列索引不属于a的剩余元素

实现代码（JavaScript） #

// 主矩阵
const w = [
  [0, 0.6, 0.3, 0.3, 0.5, 0.5],
  [0.9, 0, 0.7, 0.8, 0.2, 0.5],
  [0.2, 0.8, 0, 0.4, 0.6, 0.3],
  [0.1, 0.7, 0.6, 0, 0.7, 0.2],
  [0.9, 0.8, 0.5, 0.3, 0, 0.7],
  [0.8, 0.6, 0.3, 0.7, 0.7, 0],
];

function selection(a) {
  const columnSet = new Set(a);
  const subMatrixA = [];
  const subMatrixB = [];

  // 遍历选中的行
  for (const rowIdx of a) {
    const currentRow = w[rowIdx];
    const rowA = [];
    const rowB = [];

    // 拆分当前行的列元素
    for (let colIdx = 0; colIdx < currentRow.length; colIdx++) {
      if (columnSet.has(colIdx)) {
        rowA.push(currentRow[colIdx]);
      } else {
        rowB.push(currentRow[colIdx]);
      }
    }

    subMatrixA.push(rowA);
    subMatrixB.push(rowB);
  }

  return { subMatrixA, subMatrixB };
}

// 调用示例
const { subMatrixA, subMatrixB } = selection([0, 1]);
console.log(subMatrixA); // 输出预期的subMatrixA
console.log(subMatrixB); // 输出预期的subMatrixB

代码逻辑说明 #

1. 快速列判断：将a转为Set，让列索引的存在性判断从O(n)优化为O(1)，提升效率
2. 行筛选：遍历a中的行索引，取出主矩阵中对应的行
3. 列拆分：对每一行的列元素逐一判断，根据列索引是否在a中，分别存入对应子矩阵的行数组
4. 结果返回：将两个子矩阵打包返回，方便调用方获取

验证结果 #

调用selection([0,1])后，输出与预期完全一致：

• subMatrixA：[[0, 0.6], [0.9, 0]]
• subMatrixB：[[0.3, 0.3, 0.5, 0.5], [0.7, 0.8, 0.2, 0.5]]

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