CODEWITHSUNDEEP
pythonnumpy
Alex T
Alex T

Reputation: 3754

Python function to find indexes of 1s in binary array and

I have an array which looks like this

[1, 0, 1 , 0 , 0, 1]

And I want to get those indexes that have 1 in it. So here I would get an array of [0, 2 , 5] and then based on it I would create a new array that takes these numbers and exponantiate 2 with them So the end array is

[2**0, 2**2, 2**5]

Is there a way to write it as shortly as possible?

Upvotes: 5

Views: 4629

Answers (3)

Paul Panzer
Paul Panzer

Reputation: 53029

There is np.where, np.argwhere, np.nonzero and np.flatnonzero. Take your pick.

np.where and np.nonzero are as far as I know equivalent, only there is also a three argument version of where that does something different. np.argwhere is the transpose and np.flatnonzero gives flat indices.

Upvotes: 3

Divakar
Divakar

Reputation: 221514

Here's one compact way -

2**np.where(a)[0]

Sample run -

In [83]: a = np.array([1, 0, 1 , 0 , 0, 1])

In [84]: 2**np.where(a)[0]
Out[84]: array([ 1,  4, 32])

Upvotes: 4

Reblochon Masque
Reblochon Masque

Reputation: 36662

you could use enumerate in a list comprehension:

a = [1, 0, 1 , 0 , 0, 1]
b = [2**idx for idx, v in enumerate(a) if v]
b

output:

[1, 4, 32]

Upvotes: 5

Related Questions