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文档中称为高级索引。

我可以对这个二维数组索引进行矢量化，其中第二维依赖于第一维的值吗？

1 个答案: