nickpick nickpick - 1 month ago 15
Python Question

np.choose not giving desired result after broadcasting

I would like to pick the nth elements as specified in maxsuit from suitCounts. I did broadcast the maxsuit array so I do get a result, but not the desired one. Any suggestions what I'm doing conceptually wrong is appreciated. I don't understand the result of

np.choose(self.maxsuit[:,:,None]-1, self.suitCounts)
, which is not what I'm looking for.

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


The desired result would be:

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

Answer

You could use advanced-indexing for a broadcasted way to index into the array, like so -

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

A bit cleaner way with np.ogrid to use open range arrays -

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

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