Given numpy row containing numbers from range(n)
,
I want to apply the following transformation:
[1 0 1 2] --> [[0 1 0] [1 1 0] [1 2 0] [1 2 1]]
We just go through the input list and bincount all elements to the left of the current (including).
import numpy as np
n = 3
a = np.array([1, 0, 1, 2])
out = []
for i in range(a.shape[0]):
out.append(np.bincount(a[:i+1], minlength=n))
out = np.array(out)
Is there any way to speed this up? I'm wondering if it's possible to get rid of that loop completely and use matrix magic only.
EDIT: Thanks, lbragile, for mentioning list comprehensions. It's not what I meant. (I'm not sure if it's even significant asymptotically). I was thinking about some more complex things such as rewriting this based on how bincount operation works under the hood.
You can use cumsum like so:
idx = [1,0,1,2]
np.identity(np.max(idx)+1,int)[idx].cumsum(0)
# array([[0, 1, 0],
# [1, 1, 0],
# [1, 2, 0],
# [1, 2, 1]])