I am building a function that I would like to creates a random 20x20 array consisting of the values 0, 1 and 2. I would secondly like to iterate through the array and keep a count of how many of each number are in the array. Here is my code:
%matplotlib inline import matplotlib.pyplot as plt import numpy as np import random def my_array(): rand_array = np.random.randint(0,3,(20,20)) zeros = 0 ones = 0 twos = 0 for element in rand_array: if element == 0: zeros += 1 elif element == 1: ones += 1 else: twos += 1 return rand_array,zeros,ones,twos print(my_array())
When I eliminate the for loop to try and iterate the array it works fine and prints the array however as is, the code gives this error message:
ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()
When you iterate on a multi-dimensional
numpy array, you're only iterating over the first dimension. In your example, your
element values will be 1-dimensional arrays too!
You could solve the issue with another
for loop over the values of the 1-dimensional array, but in
numpy code, using
for loops is very often a bad idea. You usually want to be using vector operations and operations broadcast across the whole array instead.
In your example, you could do:
rand_array = np.random.randint(0,3,(20,20)) # no loop needed zeros = np.sum(rand_array == 0) ones = np.sum(rand_array == 1) twos = np.sum(rand_array == 2)
== operator is broadcast over the whole array producing an boolean array. Then the
sum adds up the
True values (
True is equal to
1 in Python) to get a count.