CODEWITHSUNDEEP
pythonnumpy
Alexander Golys
Alexander Golys

Reputation: 703

Randomly choose index based on condition in numpy

Let's say I have 2D numpy array with 0 and 1 as values. I want to randomly pick an index that contains 1. Is there efficient way to do this using numpy?

I achieved it in pure python, but it's too slow.

Example input:

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

output:

(0, 1)

EDIT: For clarification: I want my function to get 2D numpy array with values belonging to {0, 1}. I want the output to be a tuple (2D index) of randomly (uniformly) picked value from the given array that is equal to 1.

EDIT2: Using Paul H's suggestion, I came up with this:

    nonzero = np.nonzero(a)
    return random.choice(list(zip(nonzero)))

But it doesn't work with numpy's random choice, only with python's. Is there a way to optimise it better?

Upvotes: 1

Views: 958

Answers (2)

Quang Hoang
Quang Hoang

Reputation: 150745

It's easier to get all the non-zero coordinates and sample from there:

xs,ys = np.where([[0, 1], [1, 0]])

# randomly pick a number:
idx = np.random.choice(np.arange(len(xs)) )

# output:
out = xs[idx], ys[idx]

Upvotes: 4

Andy L.
Andy L.

Reputation: 25239

You may try argwhere and permutation

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

b = np.argwhere(a)
tuple(np.random.permutation(b)[0])

Upvotes: 1

Related Questions