Scipy.linalg.block_diag是否有反向操作？如何高效实现？

Scipy提供了scipy.linalg.block_diag函数，能方便地将数组列表转换为块对角矩阵，示例如下：

>>> from scipy.linalg import block_diag
>>> A = [[1, 0],
         [0, 1]]
>>> B = [[3, 4, 5],
         [6, 7, 8]]
>>> C = [[7]]
>>> D = block_diag(A, B, C)
>>> D
array([[1, 0, 0, 0, 0, 0],
       [0, 1, 0, 0, 0, 0],
       [0, 0, 3, 4, 5, 0],
       [0, 0, 6, 7, 8, 0],
       [0, 0, 0, 0, 0, 7]])

问题 #

是否存在该操作的反向实现？即输入块对角矩阵和块大小列表，将其分解为数组列表，例如：

a, b, c = foo(D, block_sizes=[(2,2), (2,3), (1,1)])

如果没有内置的便捷实现，有没有比以下朴素循环更高效的写法？

def foo(matrix, block_sizes):
    result = []
    curr_row, curr_col = 0, 0 
    for nrows, ncols in block_sizes:
        result.append(matrix[curr_row:curr_row + nrows, curr_col:curr_col + ncols])
        curr_row += nrows
        curr_col += ncols
    return result

回答 #

1. 是否有内置反向函数
   目前Scipy并没有提供block_diag对应的内置反向操作函数，需要自行实现。

2. 更高效的实现方式
   你给出的朴素循环实现其实已经非常高效了——因为NumPy的切片操作返回的是原矩阵的视图（不会复制数据），循环只是依次生成这些视图，几乎没有额外开销。

如果想简化写法、避免显式的行/列计数器，可以通过计算分割点的方式实现，本质和循环效率一致：

import numpy as np

def block_diag_split(matrix, block_sizes):
    # 计算行和列的分割位置
    row_splits = np.cumsum([0] + [nrows for nrows, _ in block_sizes])
    col_splits = np.cumsum([0] + [ncols for _, ncols in block_sizes])
    
    # 遍历分割点提取块
    blocks = []
    for r_start, r_end, c_start, c_end in zip(row_splits[:-1], row_splits[1:], col_splits[:-1], col_splits[1:]):
        blocks.append(matrix[r_start:r_end, c_start:c_end])
    return blocks

或者用更简洁的列表推导式：

def block_diag_split(matrix, block_sizes):
    row_splits = np.cumsum([0] + [nr for nr, _ in block_sizes])
    col_splits = np.cumsum([0] + [nc for _, nc in block_sizes])
    return [matrix[r1:r2, c1:c2] for r1, r2, c1, c2 in zip(row_splits[:-1], row_splits[1:], col_splits[:-1], col_splits[1:])]

需要注意的是，以上方法不会校验输入矩阵是否为严格的块对角矩阵（即非块区域是否有非零值），如果需要校验，可额外添加检查步骤，但会增加一定的计算开销。

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