support@lune.dev

Want to use over 200+ MCP servers inside your coding tools like Cursor?

Asked 1 month ago by NebularAdventurer082

How can I extract whole subarrays from a 3D NumPy array based on a condition on the third dimension?

The post content has been automatically edited by the Moderator Agent for consistency and clarity.

I have a 3D NumPy array with shape (3,2,2) and I want to select entire subarrays (the complete 2D arrays along axis 1) when at least one element along the third dimension meets a condition. For example, given the array below, I want to return the full corresponding subarray if any element in the 2D child array satisfies the condition.

PYTHON
import numpy as np

arr = np.array([
 [[31.,  1.], [41.,  1.]],
 [[63.,  1.], [73.,  3.]],
 [[95.,   1.], [100., 1]]
 ])

ref = arr[(arr[:,:,0] > 41.) & (arr[:,:,0] <= 63)]
print(ref)

This returns:

PYTHON
[[63.  1.]]

However, the expected result is:

PYTHON
[[63.,  1.],[73.,  3.]]

Here, if the input value is 63, I want to retrieve the entire parent subarray without reshaping, even though I don’t know in advance if another value (like 73) exists in that group.

Another example is:

PYTHON
ref = arr[(arr[:,:,0] <= 63)]

Which currently returns:

PYTHON
[[31.  1.]
 [41.  1.]
 [63.  1.]]

But it should return the entire parent subarrays:

PYTHON
[[[31.  1.]
 [41.  1.]]
 [[63.  1.]
 [73.  1.]]]

I'm looking for a way to perform boolean indexing so that the condition is applied to determine whether to include the whole subarray (the parent) rather than just selecting individual elements.

Comments 0

•

Answers 4

•

4 Answers

Answer by StackOverflow TychoAI Agent • 1 month ago

You need to decide which of the axis‑0 subarrays (the “parent” arrays) to keep if any element along the second dimension satisfies your condition on the third dimension. One way to do that is to reduce the condition along axis 1 with np.any. For example, for your first case:

PYTHON
import numpy as np

arr = np.array([
    [[31.,  1.], [41.,  1.]],
    [[63.,  1.], [73.,  3.]],
    [[95.,  1.], [100., 1]]
])

# Create a mask on axis 1: For each 2x2 subarray, check if any element in column 0
# satisfies (value > 41) & (value <= 63)
mask = np.any((arr[:,:,0] > 41) & (arr[:,:,0] <= 63), axis=1)
ref = arr[mask]
print(ref)

This produces:

PYTHON
[[[63.  1.]
  [73.  3.]]]

Similarly, for your second example where the condition is (arr[:,:,0] <= 63):

PYTHON
mask = np.any(arr[:,:,0] <= 63, axis=1)
ref = arr[mask]
print(ref)

This returns a (2,2,2) subarray (the “parent” arrays where at least one element in the first column is ≤ 63):

PYTHON
[[[31.  1.]
  [41.  1.]]
 [[63.  1.]
  [73.  3.]]]

In summary, by reducing the condition with np.any on axis=1 you preserve the whole subarray when at least one element meets your criterion.

Tycho is an AI agent, that grounds responses in various sources like documentation, code bases, live discussions, and relevant posts. Want to chat privately with Tycho?

No comments yet.

Answer by MeteoricNomad820 • 1 month ago

I think you want

PYTHON
arr[((arr[:,:,0]>41)&(arr[:,:,0]<=63)).any(axis=1)]

and

PYTHON
arr[(arr[:,:,0] <= 63).any(axis=1)]

Some explanation. First of all, a 3D array, is also a 2D array of 1D array, or a 1D array of 2D array.

So, if the expected answer is an array of "whole parent", that is an array of 2D arrays (so a 3D array, with only some subarrays in it), you need a 1D array of booleans as index. Such as [True, False, False] to select only the 1st row, [[[31., 1.], [41., 1.]]], which is a bit the same as arr[[0]].

arr[:,:,0]>41 is a 2D array of booleans. And therefore would pick individually some pairs (that is select elements along the 1st two axis). For example arr[[[True, False], [False, False], [False, False]]] selects only the 1st pair of the 1st subarray, [[31,1]]. A bit like arr[[0],[0]] would do.

So, since you want something like [[[31., 1.], [41., 1.]]], not something like [[31,1]], we need to produce a 1D array of booleans, telling for each line (each subarray along axis 0) whether we want it or not.

Now, the comments were about how to decide whether we want a subarray or not.

If we start from

PYTHON
ref=arr[:,:,0]<=63

That is ref =

PYTHON
array([[ True,  True],
       [ True, False],
       [False, False]])

getting arr[ref] would select 3 pairs, which is not what you want (again, we don't want a 2D array of booleans as selector, since we want whole parents).

Your attempt answer, is to use ref[:,0], which is a 1D-array of 3 booleans (the first of each row) : [True, True, False], which would select the 3 first rows.

This answer, is to use ref[:,0], which is also a 1D-array of 3 booleans, each True iff one boolean at least of the row is True. So, also [True, True, False]

Difference between our two answers shows with another example, used in comment.

PYTHON
ref=arr[:,:,0]>97

PYTHON
array([[False, False],
       [False, False],
       [False,  True]])

if we use ref[:,0], that is [False, False, False], then answer is an empty array. Even the last row is not selected, tho it contains a value over 97. But we are only interested (if we say ref[:,0]) in rows whose 1st value of 1st pair is > 97

If we use ref.any(axis=1), as in this answer, that is [False, False, True], we get the last row. Because this means that we are interested in any row whose at least one pair has a 1st value>97.

We could also select rows whose only second pair has a 1st value>97 (arr[ref[:,1]]). Or rows whose one pair, but not both, has a 1st value>97 (arr[ref[:,0]^ref[:,1]]). Etc. Everything is possible. Point is, if we want to get a list of whole rows (subarrays along axis 0), then we need to build a 1D array of 3 booleans, deciding for each row if we want it all (True) or nothing (false)

No comments yet.

Answer by AuroraOrbiter123 • 1 month ago

Finally found it. The boolean index needs to be applied on the right dimension.

Assigning the indexes to a var for clarity

PYTHON
ref = (arr[:,:,0] > 41) & (arr[:,:,0] <= 63)
print(ref)
print(arr[ref[:,0]])

Result

PYTHON
[[False False]
 [ True False]
 [False False]]
[[[63.  1.]
  [73.  3.]]]

UPDATE:
As mentioned in the accepted answer .any(axis=1) is the right option
in my case since this test should also match

PYTHON
ref = (arr[:,:,0] > 97).any(axis=1)
print(arr[ref])

Result

PYTHON
[[[ 95.   1.]
  [100.   1.]]]

No comments yet.

Answer by CosmicGuide485 • 1 month ago

While using the condition directly like condition = (arr[:,:,0] > 41.) & (arr[:,:,0] <= 63), It flattens the output into a single dimension array. The flattened output will be [False, False, True, True, False, False] in this case. In the flattened output, we lose 2nd dimension structure.

Instead we should use mask for filtering.

PYTHON
import numpy as np

arr = np.array([
    [[31.,  1.], [41.,  1.]],
    [[63.,  1.], [73.,  3.]],
    [[95.,  1.], [100., 1.]]
])

mask = (arr[:, :, 0] > 41.) & (arr[:, :, 0] <= 63)

ref = arr[mask.any(axis=1)]
print(ref)

Output:

PYTHON
[[[63.  1.]
  [73.  3.]]]

No comments yet.

Discussion

No comments yet.

How can I extract whole subarrays from a 3D NumPy array based on a condition on the third dimension?

4 Answers

Discussion

Similar Posts

How to extract complete 2D subarrays from a 3D NumPy array using boolean conditions on inner elements?