Boolean Masking
Problem Statement
Use boolean conditions to filter NumPy arrays.
Background
Boolean indexing selects elements based on conditions:
arr[arr > 5] # All elements greater than 5
arr[arr % 2 == 0] # All even elements
Your Task
Write a function filter_array(arr, threshold) that returns a dictionary with:
- "above": Elements strictly greater than threshold (as list)
- "below_eq": Elements less than or equal to threshold (as list)
- "count_above": Count of elements above threshold
Example:
arr = [1, 5, 3, 8, 2, 9, 4, 7], threshold = 5
{'above': [8, 9, 7], 'below_eq': [1, 5, 3, 2, 4], 'count_above': 3}arr[arr > 5] gives [8, 9, 7], arr[arr <= 5] gives [1, 5, 3, 2, 4]
Constraints:
- Use boolean indexing arr[condition]
- threshold will be a number
More from NumPy Foundations
1. Background Knowledge
NumPy Arrays and Vectorized Operations NumPy arrays enable efficient, vectorized computations without explicit loops, leveraging optimized C-level code for speed. Boolean indexing is a core feature: applying a condition (e.g., arr > threshold) to an array returns a boolean mask—a same-shaped array of True/False values indicating where the condition holds. This mask acts as an index: arr[boolean_mask] selects only elements where the mask is True, returning a 1D array of matching values.
Boolean Masking Mechanics Masks support element-wise comparisons (>, <=, ==) broadcast across the array. For example, arr > 5 creates a mask; slicing with it (arr[arr > 5]) yields filtered values as a new array (a copy, not a view). Counts use np.sum(mask) since True is 1 and False is 0 in numeric contexts. This avoids Python loops, making it scalable for large datasets in data analysis.
2. Algorithm/Approach
Use boolean masks to partition the array into subsets based on the threshold:
- Generate masks for above (arr > threshold) and below_eq (arr <= threshold).
- Apply masks to extract values: convert filtered arrays to lists.
- Compute count via summing the above mask. This filter-and-collect pattern exploits NumPy's vectorization for O(n) time, where n is array length.
3. Step-by-Step Strategy
- Validate inputs: Ensure arr is a NumPy array (use np.asarray if needed).
- Create masks:
- above_mask = arr > threshold
- below_eq_mask = arr <= threshold (or ~above_mask for efficiency).
- Extract values:
- above_values = arr[above_mask].tolist()
- below_eq_values = arr[below_eq_mask].tolist()
- Count above: count_above = np.sum(above_mask) (or len(above_values)).
- Build dictionary: return {"above": above_values, "below_eq": below_eq_values, "count_above": count_above}.
4. Common Pitfalls
- List conversion: NumPy arrays must be .tolist() for dictionary values; raw arrays fail sample output.
- Strict inequality: Use > (not >=) for "strictly greater" in above.
- Input types: threshold may be scalar/float; NumPy auto-broadcasts, but test edge cases (e.g., empty arrays, all below).
- Views vs. copies: Boolean indexing returns copies—safe here, but avoid mutating if sharing data.
- Non-NumPy inputs: Function assumes arr is array-like; wrap with np.array(arr) for robustness.
5. Time & Space Complexity
- Time: O(n) for mask creation, filtering, and counting—all vectorized single passes.
- Space: O(n) for masks + O(k) for output lists (k ≤ n); temporary, as masks are boolean (1 byte/element). Scales linearly; optimal for NumPy.