Numpy: Pick elements based on bool array

I've got an array and a boolean array (as one hot encoding)

a = np.arange(12).reshape(4,3)
b = np.array([
    [1,0,0],
    [0,1,0],
    [0,0,1],
    [0,0,1],
], dtype=bool)

print(a)
print(b)
# [[ 0  1  2]
#  [ 3  4  5]
#  [ 6  7  8]
#  [ 9 10 11]]
# [[ True False False]
#  [False  True False]
#  [False False  True]
#  [False False  True]]

And I would like to pick elements using a boolean array

print(a[:, [True, False, False]])
# array([[0],
#        [3],
#        [6],
#        [9]])

print(a[:, [False, True, False]])
# array([[ 1],
#        [ 4],
#        [ 7],
#        [10]])

But this picks based on the same template boolean for all rows. I would like to perform this on a per row basis:

print(a[:, b])
# IndexError: too many indices for array

What should I put in ... so I get:

print(a[:, ...])
# array([[0],
#        [4],
#        [8],
#        [11]])

EDIT: This is analogous to what was used in the infamous CS231 course:

dscores = a
num_examples = 4 
# They had 300
y = b
dscores[range(num_examples),y]
# equivalent to
# a{:,b]

EDIT 2: In CS231 example, y is one dimensional and is not one hot encoded!

They were doing dscores[[rowIdx],[columnIdx]]

Upvotes: 1

Views: 828

Answers (2)

kmario23
kmario23

Reputation: 61305

Here is an alternative way of doing the same. Please note that this is inefficient when compared to advanced indexing. It's just for pedagogical purposes and to illustrate that a problem can be solved using more than one approach.

In [275]: np.add.reduce(a*b, axis=1, keepdims=True)
Out[275]: 
array([[ 0],
       [ 4],
       [ 8],
       [11]])

Upvotes: 0

BENY
BENY

Reputation: 323226

After filter by b broadcast it

a[b][:,None]
Out[168]: 
array([[ 0],
       [ 4],
       [ 8],
       [11]])

Or

a[b,None]
Out[174]: 
array([[ 0],
       [ 4],
       [ 8],
       [11]])

Upvotes: 3

Related Questions