如何根据起始点列表修剪NumPy数组（避免使用for循环）

现有一个以np.array表示的表格A：

A = [[12,412,42,54],
     [144,2,42,4],
     [2,43,22,10]]

另有列表L，包含A中每行的期望起始点：

L=[0,2,1]

期望得到输出数组B：

B = [[12,412,42,54],
     [42,4,np.nan,np.nan],
     [43,22,10,np.nan]]

注：希望避免使用for循环。

利用NumPy的广播机制可以实现无循环操作，具体代码如下：

import numpy as np

A = np.array([[12,412,42,54],
              [144,2,42,4],
              [2,43,22,10]])
L = np.array([0,2,1])

rows, cols = A.shape
# 生成列索引数组，用于和每行起始点做比较
col_indices = np.arange(cols)
# 通过广播生成布尔掩码：标记每行中需要保留的列位置
mask = col_indices >= L[:, np.newaxis]
# 初始化结果数组为全nan，再填充符合条件的元素
B = np.full_like(A, np.nan, dtype=np.float64)
B[mask] = A[mask]

print(B)

思路说明 #

1. 将原始列表转为NumPy数组，便于执行向量化操作
2. 生成列索引数组col_indices，通过L[:, np.newaxis]将起始点列表转为列向量，借助广播特性逐行判断列索引是否大于等于对应行的起始点，得到布尔掩码
3. 创建与A形状一致的全nan数组，最后将掩码为True的位置替换为A中对应元素，即可得到目标结果

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