如何在NumPy中将矩阵右上三角值移至右下区域？

给定如下右上三角矩阵，需要将其中的非NaN值按示例移至矩阵底部：

import numpy as np

input_array = np.array([
    [np.nan,      1,      2,      3],
    [np.nan, np.nan,      4,      5],
    [np.nan, np.nan, np.nan,      6],
    [np.nan, np.nan, np.nan, np.nan],
])

# 期望输出
result = np.array([
    [np.nan, np.nan, np.nan, np.nan],
    [np.nan, np.nan, np.nan,      3],
    [np.nan, np.nan,      2,      5],
    [np.nan,      1,      4,      6],
])

我的meshgrid实现方案 #

import numpy as np

n = input_array.shape[0]
j, i = np.meshgrid(np.arange(n), np.arange(n))
i2 = i - np.arange(1,n+1)[::-1]
result = input_array[i2, j]

更简洁高效的实现思路 #

可以利用布尔索引直接定位源元素和目标位置，避免创建冗余的网格数组，代码更紧凑且运行效率更高：

import numpy as np

n = input_array.shape[0]
result = np.full_like(input_array, np.nan)
# 生成标记右上三角元素的布尔矩阵
src_mask = np.arange(n)[:, None] < np.arange(n)
# 生成标记结果矩阵中目标位置的布尔矩阵
dst_mask = np.arange(n)[:, None] >= (n - np.arange(n))
# 直接批量赋值
result[dst_mask] = input_array[src_mask]

或者用三角索引的方式，更直观地定位元素位置：

import numpy as np

n = input_array.shape[0]
result = np.full_like(input_array, np.nan)
# 获取右上三角的(i,j)索引（i < j）
i_src, j_src = np.triu_indices(n, k=1)
# 计算目标行索引：每列j的元素放在从n-j开始的连续行中
i_dst = (n - j_src) + i_src
# 赋值
result[i_dst, j_src] = input_array[i_src, j_src]

这两种方法都不需要创建meshgrid生成的二维索引数组，内存占用更低，向量化操作的执行效率也更优，同时代码逻辑更清晰易懂。

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