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]]]