3d numpy阵列的改造

Question

我有 3d 阵列，我需要将其右侧部分设置为零。 对于数组的每个二维切片 (n, :, :)，列的索引应取自向量 b。 该索引定义了分离点——左右部分，如下图所示。

a_before = [[[ 1  2  3  4]
             [ 5  6  7  8]
             [ 9 10 11 12]
             [13 14 15 16]]

            [[17 18 19 20]
             [21 22 23 24]
             [25 26 27 28]
             [29 30 31 32]]

            [[33 34 35 36]
             [37 38 39 40]
             [41 42 43 44]
             [45 46 47 48]]]

a_before.shape = (3, 4, 4)
b = (2, 3, 1)

a_after_1 = [[[ 1  2  0  0]
              [ 5  6  0  0]
              [ 9 10 0 0]
              [13 14 0 0]]

             [[17 18 19 0]
              [21 22 23 0]
              [25 26 27 0]
              [29 30 31 0]]

             [[33 0 0 0]
              [37 0 0 0]
              [41 0 0 0]
              [45 0 0 0]]]

在此之后，对于每个二维切片 (n, :, :)，我必须从 c 向量中获取列的索引，然后乘以从向量 d 中获取的相应值。

c = (1, 2, 0)
d = (50, 100, 150)
a_after_2 = [[[ 1  100  0  0]
              [ 5  300  0  0]
              [ 9 500 0 0]
              [13 700 0 0]]

             [[17 18 1900 0]
              [21 22 2300 0]
              [25 26 2700 0]
              [29 30 3100 0]]

             [[4950 0 0 0]
              [5550 0 0 0]
              [6150 0 0 0]
              [6750 0 0 0]]]

我做到了，但我的版本看起来很难看。 也许有人可以帮助我。

PS 我想避免 for 循环，只使用 numpy 方法。

谢谢你。

Answer 1

这是一个没有循环的版本。

In [232]: A = np.arange(1,49).reshape(3,4,4)
In [233]: b = np.array([2,3,1])
In [234]: d = np.array([50,100,150])
In [235]: I,J = np.nonzero(b[:,None]<=np.arange(4))
In [236]: A[I,:,J]=0
In [237]: A[np.arange(3),:,b-1] *= d[:,None]
In [238]: A
Out[238]: 
array([[[   1,  100,    0,    0],
        [   5,  300,    0,    0],
        [   9,  500,    0,    0],
        [  13,  700,    0,    0]],

       [[  17,   18, 1900,    0],
        [  21,   22, 2300,    0],
        [  25,   26, 2700,    0],
        [  29,   30, 3100,    0]],

       [[4950,    0,    0,    0],
        [5550,    0,    0,    0],
        [6150,    0,    0,    0],
        [6750,    0,    0,    0]]])

在我开发这个之前，我写了一个迭代版本。 它帮助我将问题可视化。

In [240]: Ac = np.arange(1,49).reshape(3,4,4)
In [241]: 
In [241]: for i,v in enumerate(b):
     ...:     Ac[i,:,v:]=0
     ...: 
In [242]: for i,(bi,di) in enumerate(zip(b,d)):
     ...:     Ac[i,:,bi-1]*=di

它可能更容易理解，从这个意义上说，不那么丑陋！

事实上，您的A具有“随心所欲”的中间维度，这使问题“矢量化”变得复杂。

使用 (3,4) 二维数组，解决方案是：

In [251]: Ab = Ac[:,0,:]
In [252]: Ab[b[:,None]<=np.arange(4)]=0
In [253]: Ab[np.arange(3),b-1]*=d

Answer 2

这里是：

import numpy as np

a = np.arange(1,49).reshape(3,4,4)
b = np.array([2,3,1])
c = np.array([1,2,0])
d = np.array([50,100,150])

for i in range(len(b)):
    a[i,:,b[i]:] = 0

for i,j in enumerate(c):
    a[i,:,j] = a[i,:,j]* d[i]

print(a)

#
[[[   1  100    0    0]
  [   5  300    0    0]
  [   9  500    0    0]
  [  13  700    0    0]]

 [[  17   18 1900    0]
  [  21   22 2300    0]
  [  25   26 2700    0]
  [  29   30 3100    0]]

 [[4950    0    0    0]
  [5550    0    0    0]
  [6150    0    0    0]
  [6750    0    0    0]]]