How to count how many times a value is in a table


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
        twos += 1

  return rand_array,zeros,ones,twos


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)

The == 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.