如何按独立位移向量对二维Numpy数组各行左移并填充NaN？

问题说明 #

给定二维Numpy数组A和一维位移向量shift，要求对A的每一行，按照shift中对应位置的数值执行左移操作，右侧空缺的位置用NaN填充。示例如下：

import numpy as np

# 输入数组
A = np.array([[4.0, 3, 2],
              [1, 2, 3],
              [0, -1, 5]])
shift = np.array([1,2,1])

# 期望输出
expected_out = np.array([[3, 2, np.nan],
                         [3, np.nan, np.nan],
                         [-1, 5, np.nan]])

实现方案 #

方案1：Numpy广播索引（高效无循环）

利用Numpy的广播机制生成索引矩阵，一次性完成所有行的赋值，适合处理大规模数组：

import numpy as np

A = np.array([[4.0, 3, 2],
              [1, 2, 3],
              [0, -1, 5]])
shift = np.array([1,2,1])

rows, cols = A.shape
# 为每一行生成目标位置对应的原数组列索引
col_idx = np.arange(cols) - shift[:, np.newaxis]
# 标记哪些索引是有效的（在原数组列范围内）
valid_mask = col_idx >= 0

# 初始化输出数组，默认填充NaN
output = np.full_like(A, np.nan)
# 将有效索引对应的元素赋值到输出数组
output[valid_mask] = A[np.broadcast_to(np.arange(rows)[:, None], (rows, cols))[valid_mask], col_idx[valid_mask]]

print(output)

运行后输出与期望结果一致。

方案2：逐行循环（简洁直观）

如果数组规模较小，逐行处理的逻辑更易懂，代码维护成本低：

import numpy as np

A = np.array([[4.0, 3, 2],
              [1, 2, 3],
              [0, -1, 5]])
shift = np.array([1,2,1])

rows, cols = A.shape
output = np.full_like(A, np.nan)

for i in range(rows):
    s = shift[i]
    # 仅当位移量小于列数时，截取对应元素填充到输出行左侧
    if s < cols:
        output[i, :cols - s] = A[i, s:]

print(output)

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