如何使用numpy将一维数组转换为行移位且末尾补零的二维数组

你可以通过以下两种方式实现需求，两种方案均适配任意长度的一维输入数组：

方案1：易读基础版（适合小数组场景） #

逻辑直观无理解门槛，代码如下：

import numpy as np

def build_shifted_square(arr):
    # 将输入转为numpy一维数组
    input_arr = np.asarray(arr).ravel()
    n = len(input_arr)
    # 初始化等边长的全零方阵
    result = np.zeros((n, n), dtype=input_arr.dtype)
    # 逐行填充对应偏移的元素
    for row_idx in range(n):
        fill_length = n - row_idx
        result[row_idx, :fill_length] = input_arr[row_idx:]
    return result

方案2：向量化高效版（适合大数组场景） #

无Python层级循环，处理长数组时性能远高于基础版，代码如下：

import numpy as np

def build_shifted_square_fast(arr):
    input_arr = np.asarray(arr).ravel()
    n = len(input_arr)
    # 生成补零后的扩展数组
    padded_arr = np.concatenate([input_arr, np.zeros(n, dtype=input_arr.dtype)])
    # 广播生成滑动窗口索引
    index_matrix = np.arange(n)[:, None] + np.arange(n)
    return padded_arr[index_matrix]

效果验证 #

测试第一个示例：

a = [1, 2, 3, 4, 5]
print(build_shifted_square(a))

输出结果：

[[1 2 3 4 5]
 [2 3 4 5 0]
 [3 4 5 0 0]
 [4 5 0 0 0]
 [5 0 0 0 0]]

测试第二个示例：

a = [1, 6, 2, 3, 8]
print(build_shifted_square_fast(a))

输出结果：

[[1 6 2 3 8]
 [6 2 3 8 0]
 [2 3 8 0 0]
 [3 8 0 0 0]
 [8 0 0 0 0]]

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