首页 > 解决方案 > 将 2-d 数组添加到 3-d 数组中,索引不断变化

问题描述

我正在尝试将 2-d 数组添加到具有不断变化的 index 的 3-d 数组中,我想出了以下代码:

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 循环。

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

谢谢你。

标签: pythonpython-3.xnumpy

解决方案


使用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是硬拷贝,如果你需要修改它是必要的。


推荐阅读