Search code examples
pythonnumpydecision-tree

numpy bincount sequential slices of array


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.


Solution

  • 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]])