CODEWITHSUNDEEP
pythonarraysnumpy
j c
j c

Reputation: 2129

Numpy array multiple mask

Trying to slice and average a numpy array multiple times, based on an integer mask array:

i.e.

import numpy as np

data = np.arange(11)
mask = np.array([0, 1, 1, 1, 0, 2, 2, 3, 3, 3, 3])

results = list()
for maskid in range(1,4):
    result = np.average(data[mask==maskid])
    results.append(result)
output = np.array(result)

Is there a way to do this faster, aka without the "for" loop?

Upvotes: 1

Views: 406

Answers (1)

Divakar
Divakar

Reputation: 221514

One approach using np.bincount -

np.bincount(mask, data)/np.bincount(mask)

Another one with np.unique for a generic case when the elements in mask aren't necessarily sequential starting from 0 -

_,ids, count = np.unique(mask, return_inverse=1, return_counts=1)
out = np.bincount(ids, data)/count

Upvotes: 1

Related Questions