python - np.choose 广播后没有给出期望的结果

Question

我想从suitCounts中选择maxsuit中指定的第n个元素。我确实广播了 maxsuit 数组，所以我确实得到了结果，但不是想要的结果。任何我在概念上做错的建议都值得赞赏。我不明白 np.choose(self.maxsuit[:,:,None]-1, self.suitCounts) 的结果，这不是我要找的。

>>> self.maxsuit
Out[38]: 
array([[3, 3],
       [1, 1],
       [1, 1]], dtype=int64)

>>> self.maxsuit[:,:,None]-1
Out[33]: 
array([[[2],
        [2]],

       [[0],
        [0]],

       [[0],
        [0]]], dtype=int64)
>>> self.suitCounts
Out[34]: 
array([[[2, 1, 3, 0],
        [1, 0, 3, 0]],

       [[4, 1, 2, 0],
        [3, 0, 3, 0]],

       [[2, 2, 0, 0],
        [1, 1, 1, 0]]])
>>> np.choose(self.maxsuit[:,:,None]-1, self.suitCounts)
Out[35]: 
array([[[2, 2, 0, 0],
        [1, 1, 1, 0]],

       [[2, 1, 3, 0],
        [1, 0, 3, 0]],

       [[2, 1, 3, 0],
        [1, 0, 3, 0]]])

期望的结果是:

[[3,3],[4,3],[2,1]]

Answer 1

您可以使用advanced-indexing对于数组索引的广播方式，就像这样 -

In [415]: val     # Data array
Out[415]: 
array([[[2, 1, 3, 0],
        [1, 0, 3, 0]],

       [[4, 1, 2, 0],
        [3, 0, 3, 0]],

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

In [416]: idx     # Indexing array
Out[416]: 
array([[3, 3],
       [1, 1],
       [1, 1]])

In [417]: m,n = val.shape[:2]

In [418]: val[np.arange(m)[:,None],np.arange(n),idx-1]
Out[418]: 
array([[3, 3],
       [4, 3],
       [2, 1]])

使用 np.ogrid 更干净一点使用开放范围数组 -

In [424]: d0,d1 = np.ogrid[:m,:n]

In [425]: val[d0,d1,idx-1]
Out[425]: 
array([[3, 3],
       [4, 3],
       [2, 1]])