CODEWITHSUNDEEP
pythonarraysnumpymultidimensional-array
ssast
ssast

Reputation: 789

numpy count elements across axis 0 matching values from another array

Given a 3D array such as:

array = np.random.randint(1, 6, (3, 3, 3))

and an array of maximum values across axis 0:

max_array = array.max(axis=0)

Is there a vectorised way to count the number of elements in axis 0 of array which are equal to the value of the matching index in max_array? For example, if array contains [1, 3, 3] in one axis 0 position, the output is 2, and so on for the other 8 positions, returning an array with the counts.

Upvotes: 1

Views: 1975

Answers (1)

unutbu
unutbu

Reputation: 880677

To count the number of values in x which equal the corresponding value in xmax, you could use:

(x == xmax).sum(axis=0)

Note that since x has shape (3,3,3) and xmax has shape (3,3), the expression x == xmax causes NumPy to broadcast xmax up to shape (3,3,3) where the new axis is added on the left.

For example,

import numpy as np
np.random.seed(2015)
x = np.random.randint(1, 6, (3,3,3))
print(x)
# [[[3 5 5]
#   [3 2 1]
#   [3 4 1]]

#  [[1 5 4]
#   [1 4 1]
#   [2 3 4]]

#  [[2 3 3]
#   [2 1 1]
#   [5 1 2]]]

xmax = x.max(axis=0)
print(xmax)
# [[3 5 5]
#  [3 4 1]
#  [5 4 4]]

count = (x == xmax).sum(axis=0)
print(count)
# [[1 2 1]
#  [1 1 3]
#  [1 1 1]]

Upvotes: 2

Related Questions