python - 使用 numpy 根据数组的条件索引创建矩阵

Question

我想根据 n + m 长度数组的元素创建一个 n by m 矩阵。

这里有一个简单的 double for 循环就足够了，但我希望有一个权宜之计的解决方案。矩阵会比较小。

n = 4
m = 6
s = n + m

array = np.arange(s)  # note: arange is only for example. real array varies.
matrix = np.zeros((n,m))

for i in range(n):
    for j in range (m):
        matrix[i,j] = array[i+j]

我发现理解比双循环更快

matrix3 = [[array[i+j] for i in range(m)] for j in range(n)]

有没有更快的方法？

另一个好处是合并模运算符。我实际上只需要 i+j % 2 == 0 的索引。在双 for 循环中，取模方法似乎要快一点，但这对于通过 numpy 生成此矩阵可能不方便或权宜之计。

不这样做也没关系，因为矩阵乘法会发生在之后，必要的元素无论如何都会乘以零。我提到模数只是在这种情况下，这会导致更快的解决方案。

对于这个 MWE

for i in range(n):
    for j in range (m):
        if (i + j) % 2 == 0:
            matrix[i,j] = array[i+j]

注意:

我要求一个 numpy 解决方案，假设 numpy 最快，但任何纯 python(包括 numpy/scipy)解决方案都可以，只要它比纯 python double for 循环更快

动机:

我正在尝试从双 for 循环中删除对数组的所有依赖项，以便我可以使用广播而不是双 for 循环。这是剩下的最后一个数组

Answer 1

您可以在 array 中使用高级索引。为了提高效率，您可以将模板数组中已经存在的奇数位置归零。

np.where(np.arange(m+n)&1,0,array)[sum(np.ogrid[:n,:m])]
# array([[0, 0, 2, 0, 4, 0],
#        [0, 2, 0, 4, 0, 6],
#        [2, 0, 4, 0, 6, 0],
#        [0, 4, 0, 6, 0, 8]])

或(更快)

template = np.where(np.arange(m+n)&1,0,array)
np.lib.stride_tricks.as_strided(template,(n,m),2*template.strides)

这是一个“压缩” View ，如果您需要修改条目，您必须制作一个副本(它仍然会更快)。