Johannes Wiesner
Johannes Wiesner

Reputation: 1307

Get indices from one randomly chosen true element in a boolean array

I have a boolean array from which I would like the indices of one randomly chosen element that equals True. The output should be a tuple with the (x,y,z) indices of that element.

Is there a more elegant and/or efficient way to do this instead of doing the following?

import numpy as np
rng = np.random.RandomState(42)

# create a random 3D boolean array
m = rng.choice(a=[False, True],size=(3,3,3))

# get (x,y,z) from one random cell in the boolean array which equals True
indices_where_true = np.where(m)
random_int = rng.randint(len(indices_where_true[0]),size=1)
random_index = (indices_where_true[0][random_int][0],
                indices_where_true[1][random_int][0],
                indices_where_true[2][random_int][0])

Upvotes: 2

Views: 745

Answers (2)

Sebastian Baltser
Sebastian Baltser

Reputation: 758

The note in the documentation for numpy.where states:

When only condition is provided, this function is a shorthand for np.asarray(condition).nonzero(). Using nonzero directly should be preferred, as it behaves correctly for subclasses.

So you should replace where with nonzero. After sampling a random integer you can pick the coordinates by taking all elements at that index along the first axis of indices_where_true using numpy.take:

from numpy import random
rng = np.random.RandomState(42)

# create a random 3D boolean array
m = rng.choice(a=[False, True],size=(3,3,3))

# get (x,y,z) from one random cell in the boolean array which equals True
indices_where_true = np.nonzero(m)
random_int = rng.randint(len(indices_where_true[0]),size=1)
random_index = np.take(indices_where_true, random_int, axis=1)

Upvotes: 0

Quang Hoang
Quang Hoang

Reputation: 150735

Use np.argwhere instead of np.where:

true_idx = np.argwhere(m)

random_idx = rng.randint(len(true_idx),size=1)
random_index = true_idx[random_idx]
# array([[0, 1, 2]])

Upvotes: 2

Related Questions