如何提取分块对角矩阵的块外上三角元素（类似np.triu_idx）

要提取分块对角矩阵中属于上三角且不属于对角块的元素，可以通过掩码矩阵精准定位目标区域，具体实现步骤如下：

步骤说明 #

1. 记录原始对角块的尺寸，用于定位对角块在合并后矩阵中的位置
2. 生成对角块区域的掩码（标记所有属于对角块的元素位置）
3. 生成上三角区域的掩码（标记所有上三角位置，含对角线）
4. 合并两个掩码，得到「上三角且非对角块」的目标区域掩码
5. 利用掩码提取对应元素

完整代码实现 #

import numpy as np
import scipy as sp

# 定义原始对角块
A = np.array([
    [1, 1, 1],
    [1, 1, 1],
    [1, 1, 1]
])

B = np.array([
    [2, 2],
    [2, 2]
])

C = np.array([[3]])

# 生成分块对角矩阵
D = sp.linalg.block_diag(A, B, C)

# 1. 获取每个对角块的尺寸
blocks = [A, B, C]
block_sizes = [block.shape[0] for block in blocks]

# 2. 生成对角块掩码
n = D.shape[0]
block_mask = np.zeros((n, n), dtype=bool)
start_idx = 0
for size in block_sizes:
    end_idx = start_idx + size
    # 标记当前对角块的区域为True
    block_mask[start_idx:end_idx, start_idx:end_idx] = True
    start_idx = end_idx

# 3. 生成上三角掩码（含对角线）
upper_tri_mask = np.triu(np.ones_like(D, dtype=bool), k=0)

# 4. 生成目标区域掩码：上三角 且 不在对角块内
target_mask = upper_tri_mask & ~block_mask

# 5. 提取目标元素
target_elements = D[target_mask]

print("目标元素：", target_elements)
# 输出：目标元素： [0 0 0 0 0 0]

关键细节说明 #

• 如果需要提取不含对角线的上三角（仅i<j的位置），只需将上三角掩码的k参数改为1：upper_tri_mask = np.triu(np.ones_like(D, dtype=bool), k=1)
• 掩码矩阵的逻辑运算& ~用于取「上三角区域减去对角块区域」的交集，精准定位目标位置

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