CODEWITHSUNDEEP
pythonarraysnumpyboolean-logicobject-slicing
stagermane
stagermane

Reputation: 1063

Slicing arrays based on boolean array in python

I need to slice an array of xyz coordinates based on conditions in a boolean array (where the boolean array is 1D).

If my boolean array is

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

I need it to slice it to produce the following index arrays:

[0, 1, 2, 3, 6, 7, 10, 11, 12] ([:-2] between True indices)

The ultimate desired output would be an array of the xyz coordinates at those indices:

[xyz[0], xyz[1], xyz[2] xyz[3], xyz[6] xyz[7], xyz[10] xyz[11], xyz[12]]

The other two slices (with similar desired outputs) are as follows:

[1, 2, 3, 4, 7, 8, 11, 12, 13] ([1:-1] between True indices)

[2, 3, 4, 5, 8, 9, 12, 13, 14] ([2:] between True indices)

Is there a pythonic way to do this without a list comprehension?

Thank you!

Upvotes: 0

Views: 1893

Answers (1)

Maxim Egorushkin
Maxim Egorushkin

Reputation: 136425

It looks like numpy.ndarray.nonzero is what you need:

a = np.asarray([1, 0, 0, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 0, 1])
print(a.nonzero())

Outputs:

array([ 0,  5,  6,  9, 10, 14])

Upvotes: 1

Related Questions