NumPy Conditional Selection and Element Extraction

Using np.where() to Find Indices

The simplest form of np.where() returns the indices where a condition is True. It is commonly used to find positions of elements that meet specific criteria in an array.

Example: Find Where Elements Are Greater Than 15

import numpy as np

arr = np.array([5, 17, 23, 10, 45, 3])

# Find indices where elements are greater than 15
indices = np.where(arr > 15)

print("Original array:", arr)
print("Indices where arr > 15:", indices)
print("Values:", arr[indices])

import numpy as np

arr = np.array([5, 17, 23, 10, 45, 3])

# Find indices where elements are greater than 15
indices = np.where(arr > 15)

print("Original array:", arr)
print("Indices where arr > 15:", indices)
print("Values:", arr[indices])

How It Works

• arr > 15 creates a boolean array like [False, True, True, False, True, False].
• np.where() returns the indices of the True values — in this case, [1, 2, 4].
• You can use these indices to extract matching elements: arr[indices].
• The result is returned as a tuple (e.g., (array([1, 2, 4]),)) to support multi-dimensional arrays too.

Output

Original array: [ 5 17 23 10 45  3]
Indices where arr > 15: (array([1, 2, 4]),)
Values: [17 23 45]

Original array: [ 5 17 23 10 45  3]
Indices where arr > 15: (array([1, 2, 4]),)
Values: [17 23 45]

💡 Tip: You can use this method to locate and operate on elements conditionally without writing explicit loops.

Finding Non-Zero Elements with np.nonzero()

The np.nonzero() function returns the indices of all non-zero elements in a NumPy array. It's useful when you want to locate where data exists (i.e. not zero), especially in sparse arrays or binary masks.

Example: Get Indices of Non-Zero Values

import numpy as np

arr = np.array([0, 4, 0, 7, 0, 2])

# Find indices of non-zero elements
nonzero_indices = np.nonzero(arr)

print("Original array:", arr)
print("Non-zero indices:", nonzero_indices)
print("Non-zero values:", arr[nonzero_indices])

import numpy as np

arr = np.array([0, 4, 0, 7, 0, 2])

# Find indices of non-zero elements
nonzero_indices = np.nonzero(arr)

print("Original array:", arr)
print("Non-zero indices:", nonzero_indices)
print("Non-zero values:", arr[nonzero_indices])

How It Works

• arr is scanned element by element.
• np.nonzero(arr) returns the indices where the element is not equal to zero.
• The result is a tuple (to support multi-dimensional arrays), where the first element is a 1D array of matching indices.
• You can use these indices to extract the non-zero values: arr[nonzero_indices].

Output

Original array: [0 4 0 7 0 2]
Non-zero indices: (array([1, 3, 5]),)
Non-zero values: [4 7 2]

Original array: [0 4 0 7 0 2]
Non-zero indices: (array([1, 3, 5]),)
Non-zero values: [4 7 2]

💡 Tip: np.nonzero() is especially useful when working with boolean masks or sparse data structures.

Finding Unique Elements with np.unique()

The np.unique() function returns the sorted unique elements of an array. You can also get additional information such as the counts of each unique element or the indices that reconstruct the original array.

Example: Basic Usage and Counting Unique Elements

import numpy as np

arr = np.array([2, 3, 3, 1, 5, 2, 4, 3])

# Get sorted unique elements
unique_elements = np.unique(arr)

# Get unique elements and their counts
unique_elements, counts = np.unique(arr, return_counts=True)

print("Original array:", arr)
print("Unique elements:", unique_elements)
print("Counts:", counts)

import numpy as np

arr = np.array([2, 3, 3, 1, 5, 2, 4, 3])

# Get sorted unique elements
unique_elements = np.unique(arr)

# Get unique elements and their counts
unique_elements, counts = np.unique(arr, return_counts=True)

print("Original array:", arr)
print("Unique elements:", unique_elements)
print("Counts:", counts)

How It Works

• np.unique(arr) returns a sorted array of unique values.
• Using return_counts=True also returns how many times each unique value appears.
• You can use this to analyze data distributions or remove duplicates efficiently.

Output

Original array: [2 3 3 1 5 2 4 3]
Unique elements: [1 2 3 4 5]
Counts: [1 2 3 1 1]

Original array: [2 3 3 1 5 2 4 3]
Unique elements: [1 2 3 4 5]
Counts: [1 2 3 1 1]

💡 Tip: You can also get indices to reconstruct the original array using return_inverse=True or get the first occurrence indices with return_index=True.

Extracting Elements with np.extract()

The np.extract() function allows you to extract elements from an array that satisfy a given condition. It returns a 1D array containing all elements where the condition is True.

Example: Extracting Values Based on a Condition

import numpy as np

arr = np.array([10, 15, 7, 20, 3])

# Extract elements greater than 10
extracted = np.extract(arr > 10, arr)

# Extraction using boolean indexing
extracted_bool_index = arr[arr > 5]

print("Original array:", arr)
print("Extracted elements (arr > 10):", extracted)
print("Extracted elements (arr > 5):", extracted_bool_index)

import numpy as np

arr = np.array([10, 15, 7, 20, 3])

# Extract elements greater than 10
extracted = np.extract(arr > 10, arr)

# Extraction using boolean indexing
extracted_bool_index = arr[arr > 5]

print("Original array:", arr)
print("Extracted elements (arr > 10):", extracted)
print("Extracted elements (arr > 5):", extracted_bool_index)

How It Works

• np.extract(condition, arr) returns the values in arr where the condition is True.
• It returns a 1D array regardless of the shape of the original array.
• This function is useful for quickly selecting elements that meet a criterion without manual indexing.
• You can also achieve the same result more concisely using boolean indexing: arr[condition].

Output

Original array: [10 15  7 20  3]
Extracted elements (arr > 10): [15 20]
Extracted elements (arr > 5): [10 15  7 20]

Original array: [10 15  7 20  3]
Extracted elements (arr > 10): [15 20]
Extracted elements (arr > 5): [10 15  7 20]

💡 Tip: For just indices, use np.where(). For direct values matching a condition, np.extract() is simpler.

Frequently Asked Questions