nickpick - 1 year ago 71

Python Question

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

`>>> 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 Source

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