python - 我可以向量化这个二维数组索引，其中第二维取决于第一维的值吗？

Question

在下面的示例中，我有一个 2D 数组，其中有一些经过移位和填充的实际结果。移位取决于行(填充用于根据 numpy 的要求使数组成为矩形)。是否可以在不使用 Python 循环的情况下提取真实结果？

import numpy as np

# results are 'shifted' where the shift depends on the row
shifts = np.array([0, 8, 4, 2], dtype=int)
max_shift = shifts.max()
n = len(shifts)

t = 10 # length of the real results we care about

a = np.empty((n, t + max_shift), dtype=int)
b = np.empty((n, t), dtype=int)

for i in range(n):
    a[i] = np.concatenate([[0] * shifts[i],               # shift
                           (i+1) * np.arange(1, t+1),     # real data
                           [0] * (max_shift - shifts[i])  # padding
                           ])
print "shifted and padded\n", a

# I'd like to remove this Python loop if possible
for i in range(n):
    b[i] = a[i, shifts[i]:shifts[i] + t]

print "real data\n", b

Answer 1

您可以使用两个数组来获取数据:

a[np.arange(4)[:, None], shifts[:, None] + np.arange(10)]

或者:

i, j = np.ogrid[:4, :10]
a[i, shifts[:, None]+j]

这在 NumPy 文档中称为高级索引。