Find bundles of repeating values and perform operations on them separately

user_13476168 0
Viewed 59

I have pandas Series of the type

    a = pd.Series([1,  4,3,5,  7,5,  5,6,7,6,  7,  6,2,2,  6,  9])

and a numpy array of the type

    b = np.array([0,  1,1,1,  0,0,  1,1,1,1,  0,  1,1,1,  0,  1])

Now I want to seperately generate averages of elements in a whenever there is a cluster of 1's in the corresponding parts of b, with the result

    c = pd.Series([1,  4,4,4,  7,5,  6,6,6,6,  7,  3.3,3.3,3.3  6,  9])

Does someone have an idea how to do this nicely?

pandas numpy python
2 Answers

Try with shift+cumsum , note avg of 6,2,2, is 3.333.. not 5

s = pd.Series(b,index=a.index)
a.groupby(s.ne(s.shift()).cumsum()).transform('mean').where(s.eq(1),a)

0     1.000000
1     4.000000
2     4.000000
3     4.000000
4     7.000000
5     5.000000
6     6.000000
7     6.000000
8     6.000000
9     6.000000
10    7.000000
11    3.333333
12    3.333333
13    3.333333
14    6.000000
15    9.000000
dtype: float64
user_98406370

Approach #1

Here's one NumPy way -

In [23]: ids = np.r_[0,b[:-1]!=b[1:]].cumsum()

In [24]: np.where(b==1,a.groupby(ids).transform('mean'),a)
Out[24]: 
array([1.        , 4.        , 4.        , 4.        , 7.        ,
       5.        , 6.        , 6.        , 6.        , 6.        ,
       7.        , 3.33333333, 3.33333333, 3.33333333, 6.        ,
       9.        ])

Approach #2

For performance, we can leverage np.bincount -

In [47]: v = np.bincount(ids,a)/np.bincount(ids)

In [48]: np.where(b==1,v[ids],a)
Out[48]: 
array([1.        , 4.        , 4.        , 4.        , 7.        ,
       5.        , 6.        , 6.        , 6.        , 6.        ,
       7.        , 3.33333333, 3.33333333, 3.33333333, 6.        ,
       9.        ])
user_32938810
Related