Convert a NumPy array to 0 or 1 based on threshold in Python

avatar
Borislav Hadzhiev

Last updated: Apr 10, 2024
9 min

banner

# Table of Contents

  1. Convert a NumPy array to 0 or 1 based on threshold in Python
  2. Set NumPy array elements to Zero if greater than X in Python
  3. Set the first N elements of an Array to Zero in Python

# Convert a NumPy array to 0 or 1 based on threshold in Python

Use the numpy.where() method to convert a NumPy array to 0 or 1 based on a threshold.

The numpy.where() method will set each value of the array based on whether it meets the condition.

main.py
import numpy as np arr = np.array([0.2, 0.1, 0.4, 0.7, 0.3, 0.9]) new_arr = np.where(arr > 0.5, 1, 0) print(new_arr) # ๐Ÿ‘‰๏ธ [0 0 0 1 0 1]

convert numpy array to 0 or 1 based on threshold

The code for this article is available on GitHub

Make sure you have numpy installed to be able to run the code sample.

shell
pip install numpy # ๐Ÿ‘‡๏ธ or with pip3 pip3 install numpy
We used the numpy.where() method to set all values in a numpy array that are greater than 0.5 to 1 and all values that are less than or equal to 0.5 to 0.

The numpy.where() method takes the following 3 arguments:

NameDescription
conditionIf the condition is met, return x, otherwise, return y.
xValues from which to choose if the condition is met.
yValues from which to choose if the condition is NOT met.

The numpy.where() method returns an array with elements from x where the condition is True and elements from y where the condition is False.

In other words, if the element in the array is greater than 0.5, it gets set to 1, otherwise, the element gets set to 0.

Alternatively, you can use the astype() method with a condition.

main.py
import numpy as np arr = np.array([0.2, 0.1, 0.4, 0.7, 0.3, 0.9]) new_arr = (arr > 0.5).astype(int) print(new_arr) # ๐Ÿ‘‰๏ธ [0 0 0 1 0 1]

The expression arr > 0.5 returns a boolean array containing True values for the elements that meet the condition and False for all other elements.

main.py
import numpy as np arr = np.array([0.2, 0.1, 0.4, 0.7, 0.3, 0.9]) # ๐Ÿ‘‡๏ธ [False False False True False True] print(arr > 0.5)
The code for this article is available on GitHub

The astype() method is then used to return a copy of the array casting the boolean values to integers.

main.py
import numpy as np arr = np.array([0.2, 0.1, 0.4, 0.7, 0.3, 0.9]) new_arr = (arr > 0.5).astype(int) print(new_arr) # ๐Ÿ‘‰๏ธ [0 0 0 1 0 1]
True values get converted to 1 and False values get converted to 0.

Alternatively, you can use a list comprehension.

# Convert a NumPy array to 0 or 1 based on threshold using list comprehension

This is a three-step process:

  1. Use a list comprehension to iterate over the array.
  2. Check if each item is greater than the threshold and return 0 or 1.
  3. Optionally, convert the list to an array.
main.py
import numpy as np arr = np.array([0.2, 0.1, 0.4, 0.7, 0.3, 0.9]) a_list = [ 1 if element > 0.5 else 0 for element in arr ] print(a_list) # ๐Ÿ‘‰๏ธ [0, 0, 0, 1, 0, 1] new_array = np.array(a_list) print(new_array) # ๐Ÿ‘‰๏ธ [0 0 0 1 0 1]

convert numpy array to 0 or 1 based on threshold using list comprehension

The code for this article is available on GitHub

We used a list comprehension to iterate over the numpy array.

List comprehensions are used to perform some operation for every element or select a subset of elements that meet a condition.

On each iteration, we check if the current element is greater than 0.5.

The list comprehension returns a list containing the results.

If you need to convert the list to an array, use the np.array() method.

# Set NumPy array elements to Zero if greater than X in Python

Use NumPy's indexing to set the NumPy array elements to 0 if they are greater than a certain number.

You can directly assign a zero to the slice of the array that meets the condition.

main.py
import numpy as np arr = np.array([8, 1, 2, 3, 4, 5, 6, 7]) threshold = 3 arr[arr > threshold] = 0 print(arr) # ๐Ÿ‘‰๏ธ [0 1 2 3 0 0 0 0]
The code for this article is available on GitHub

The code sample sets all array elements to 0 if they have a value greater than 3.

The conditional check arr > threshold returns an array of booleans.

main.py
import numpy as np arr = np.array([8, 1, 2, 3, 4, 5, 6, 7]) threshold = 3 # ๐Ÿ‘‡๏ธ [ True False False False True True True True] print(arr > threshold)
The array elements that are greater than 3 have a value of True and the ones that aren't, have a value of False.

You can directly use the boolean array to only set the array elements that meet the condition to 0.

main.py
import numpy as np arr = np.array([8, 1, 2, 3, 4, 5, 6, 7]) threshold = 3 arr[arr > threshold] = 0 print(arr) # ๐Ÿ‘‰๏ธ [0 1 2 3 0 0 0 0]

# Set NumPy array elements to Zero if greater than X without mutation

If you don't want to mutate the original array, use the numpy.where() method.

main.py
import numpy as np arr = np.array([8, 1, 2, 3, 4, 5, 6, 7]) threshold = 3 new_arr = np.where(arr > threshold, 0, arr) print(new_arr) # ๐Ÿ‘‰๏ธ [0 1 2 3 0 0 0 0]
The code for this article is available on GitHub

The numpy.where() method takes the following 3 arguments:

NameDescription
conditionIf the condition is met, return x, otherwise, return y.
xValues from which to choose if the condition is met.
yValues from which to choose if the condition is NOT met.

The numpy.where method returns an array with elements from x where the condition is True and elements from y where the condition is False.

In other words, if the array element is greater than the specified number, it gets set to 0, otherwise, the element is returned as is.

Alternatively, you can use a list comprehension.

# Set numpy array elements to Zero if greater than X using list comprehension

To set a NumPy array's elements to zero if they are greater than a certain threshold:

  1. Use a list comprehension to iterate over the array.
  2. Check if each item is greater than the threshold.
  3. If the condition is met, return 0, otherwise, return the element as is.
main.py
import numpy as np arr = np.array([8, 1, 2, 3, 4, 5, 6, 7]) threshold = 3 a_list = [ 0 if element > threshold else element for element in arr ] print(a_list) # ๐Ÿ‘‰๏ธ [0, 1, 2, 3, 0, 0, 0, 0] new_array = np.array(a_list) print(new_array) # ๐Ÿ‘‰๏ธ [0 1 2 3 0 0 0 0]
The code for this article is available on GitHub

We used a list comprehension to iterate over the array.

List comprehensions are used to perform some operation for every element or select a subset of elements that meet a condition.

On each iteration, we check if the current element is greater than the specified threshold.

If the condition is met, we return 0, otherwise, we return the current element.

The list comprehension returns a list containing the results.

If you need to convert the list to an array, use the np.array() method.

# Set the first N elements of an Array to Zero in Python

Use list slicing to set the first N elements of an array to zero.

main.py
import numpy as np # โœ… Set the first N elements of a Python list to 0 a_list = [8, 1, 2, 3, 4, 5, 6, 7] n = 3 a_list[:n] = [0] * n print(a_list) # ๐Ÿ‘‰๏ธ [0, 0, 0, 3, 4, 5, 6, 7] # --------------------------------------------- # โœ… Set the first N elements of NumPy array to 0 arr = np.array([8, 1, 2, 3, 4, 5, 6, 7]) n = 3 arr[:n] = 0 print(arr) # ๐Ÿ‘‰๏ธ [0 0 0 3 4 5 6 7]
The code for this article is available on GitHub

The first example shows how to set the first N elements of a native Python list to 0.

main.py
a_list = [8, 1, 2, 3, 4, 5, 6, 7] n = 3 a_list[:n] = [0] * n print(a_list) # ๐Ÿ‘‰๏ธ [0, 0, 0, 3, 4, 5, 6, 7]

We used the list-slicing assignment to set the first N elements of the list to 0.

The syntax for list slicing is my_list[start:stop:step].

The start index is inclusive and the stop index is exclusive (up to, but not including).

If the start index is omitted, it is considered to be 0, if the stop index is omitted, the slice goes to the end of the list.

Python indexes are zero-based, so the first item in a list has an index of 0, and the last item has an index of -1 or len(my_list) - 1.

The slice a_list[:n] starts at index 0 and goes up to, but not including index n.

We used the multiplication operator to create a list of N zero values on the right-hand side of the assignment.

main.py
a_list = [8, 1, 2, 3, 4, 5, 6, 7] n = 3 a_list[:n] = [0] * n print(a_list) # ๐Ÿ‘‰๏ธ [0, 0, 0, 3, 4, 5, 6, 7]

When the multiplication (*) operator is used with a list and an integer, it repeats the items in the list N times.

main.py
print([0] * 2) # ๐Ÿ‘‰๏ธ [0, 0] print([0] * 3) # ๐Ÿ‘‰๏ธ [0, 0, 0]

In its entirety, the statement replaces a slice of the first N elements of the list with a new list containing N zeros.

# Directly assigning a value to the first N array elements

If you have a numpy array, you can directly assign a zero to the first N elements of the array.

main.py
import numpy as np arr = np.array([8, 1, 2, 3, 4, 5, 6, 7]) n = 3 arr[:n] = 0 print(arr) # ๐Ÿ‘‰๏ธ [0 0 0 3 4 5 6 7]
The code for this article is available on GitHub

You can install numpy by running the following command from your terminal.

shell
pip install numpy pip3 install numpy

We can directly assign a zero to multiple elements in a numpy array.

Note that this is not a valid syntax when working with native Python lists.

The following code sample raises a TypeError exception.

main.py
a_list = [8, 1, 2, 3, 4, 5, 6, 7] n = 3 # โ›”๏ธ TypeError: can only assign an iterable a_list[:n] = 0

We are trying to assign a single 0 to a slice (multiple items) of the list.

This only works with numpy arrays, not with native Python lists.

# Replacing the first N array elements with a list of zeros

You can also replace the first N elements in the numpy array with a list containing N zeros as we did in the previous example.

main.py
import numpy as np arr = np.array([8, 1, 2, 3, 4, 5, 6, 7]) n = 3 arr[:n] = [0] * n print(arr) # ๐Ÿ‘‰๏ธ [0 0 0 3 4 5 6 7]
The code for this article is available on GitHub

The code sample sets the first N elements of the numpy array to 0 by replacing the slice with a slice containing N zero values.

You can use the tolist() method if you need to convert the numpy array to a native Python list.

main.py
import numpy as np arr = np.array([8, 1, 2, 3, 4, 5, 6, 7]) n = 3 arr[:n] = [0] * n print(arr) # ๐Ÿ‘‰๏ธ [0 0 0 3 4 5 6 7] a_list = arr.tolist() print(a_list) # ๐Ÿ‘‰๏ธ [0, 0, 0, 3, 4, 5, 6, 7]

The tolist() method converts a numpy array to a list.

# Additional Resources

You can learn more about the related topics by checking out the following tutorials:

I wrote a book in which I share everything I know about how to become a better, more efficient programmer.
book cover
You can use the search field on my Home Page to filter through all of my articles.

Copyright ยฉ 2024 Borislav Hadzhiev