快速将2维数组添加到3维数组中，且索引不断变化

Question

我正在尝试将一个二维数组添加到一个三维数组中，并且索引不断变化，我想到了以下代码：

import numpy as np

a = np.zeros([8, 3, 5])
k = 0
for i in range(2):
    for j in range(4):
        a[k, i: i + 2, j: j + 2] += np.ones([2, 2], dtype=int)
        k += 1
print(a)

这将会给我完全想要的东西：

[[[1. 1. 0. 0. 0.]
  [1. 1. 0. 0. 0.]
  [0. 0. 0. 0. 0.]]

 [[0. 1. 1. 0. 0.]
  [0. 1. 1. 0. 0.]
  [0. 0. 0. 0. 0.]]

 [[0. 0. 1. 1. 0.]
  [0. 0. 1. 1. 0.]
  [0. 0. 0. 0. 0.]]

 [[0. 0. 0. 1. 1.]
  [0. 0. 0. 1. 1.]
  [0. 0. 0. 0. 0.]]

 [[0. 0. 0. 0. 0.]
  [1. 1. 0. 0. 0.]
  [1. 1. 0. 0. 0.]]

 [[0. 0. 0. 0. 0.]
  [0. 1. 1. 0. 0.]
  [0. 1. 1. 0. 0.]]

 [[0. 0. 0. 0. 0.]
  [0. 0. 1. 1. 0.]
  [0. 0. 1. 1. 0.]]

 [[0. 0. 0. 0. 0.]
  [0. 0. 0. 1. 1.]
  [0. 0. 0. 1. 1.]]]

我希望能够更快地处理一个索引数组，因此我创建了一个数组并尝试使用np.vectorize。但正如手册所述，vectorize并不是用于提高性能的。我的目标是遍历一个形状为(10^6, 15, 15)的数组，最终需要进行10^6次迭代。我希望有一种更简洁的解决方案，可以消除所有的for循环。

这是我第一次使用stackoverflow，欢迎任何建议。

谢谢。

- Waffle

---

你是真的想要把 1 加到 a 中，还是最终想要解决其他问题，例如卷积？ - Nils Werner

@NilsWerner 我的终极目标是找到棋盘游戏中最佳的下一步。更具体地说，我想解决一个叫做1010！（类似俄罗斯方块）的游戏，因为它有三个块可以用作下一步，每个块可能有多达100个可以放置在板上的位置，所以要考虑每种可能性，可能会有多达100100100种选择。我想列出所有可能的位置和分数，并找出最佳的下一步。 - Waffle

看起来这种移动窗口可以用 as_strided 来实现，但是具体的细节会花费一些时间。我建议从 a 开始，或者对其进行简化，然后尝试找到一个索引所有 1 的方法（例如一个 (8,2,2) 的数组）。最后将该索引用于一个赋值表达式。 - hpaulj

Answer 1

一种高效的解决方案是使用 numpy.lib.stride_tricks，它可以“查看”所有可能性。

N=4 #tray size #(square)
P=3  # chunk size
R=N-P

from numpy.lib.stride_tricks import as_strided

tray = zeros((N,N),numpy.int32)
chunk = ones((P,P),numpy.int32)
tray[R:,R:] = chunk
tray = np.vstack((tray,tray))
view = as_strided(tray,shape=(R+1,R+1,N,N),strides=(4*N,4,4*N,4))
a_view = view.reshape(-1,N,N)
a_hard = a_view.copy()

这是结果：

In [3]: a_view
Out[3]: 
array([[[0, 0, 0, 0],
        [0, 1, 1, 1],
        [0, 1, 1, 1],
        [0, 1, 1, 1]],

       [[0, 0, 0, 0],
        [1, 1, 1, 0],
        [1, 1, 1, 0],
        [1, 1, 1, 0]],

       [[0, 1, 1, 1],
        [0, 1, 1, 1],
        [0, 1, 1, 1],
        [0, 0, 0, 0]],

       [[1, 1, 1, 0],
        [1, 1, 1, 0],
        [1, 1, 1, 0],
        [0, 0, 0, 0]]])

"a_view"仅是托盘上可能位置的一个视图。它不需要任何计算量，只使用了两倍的托盘空间。 "a_hard"是硬拷贝，如果您需要修改它则必须使用。

- B. M.