将2-d数组添加到3-d数组中,索引不断变化

时间:2018-08-31 15:09:39

标签: python python-3.x numpy

我试图将2维数组添加到索引不断变化的3维数组中,我想出了以下代码:

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。但是正如手册所述,矢量化不是为了性能。我的目标是遍历形状为(10 ^ 6,15,15)的数组,最后以10 ^ 6迭代。我希望有一些更清洁的解决方案可以消除所有的for-loop问题。

这是我第一次使用堆栈溢出,任何建议都会受到赞赏。

谢谢。

1 个答案:

答案 0 :(得分:2)

使用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是纸质副本,如果您需要对其进行修改,则必不可少。