sam
sam

Reputation: 19214

Calculate sum of row based on condition

I have a df

data = [[1, 10], [1, 5], [1, 4], [0, 3], [0, 3], [1, 2], [1, 1], [1, 5], [2, 8], [4, 9]]
df = pd.DataFrame(data, columns = ['no', 'count'])

which prints values as,

no  count
1    10   
1    5
1    4
0    3
0    3
1    2
1    1
1    5
2    8
4    9 

Expected output: (10-3) = 7 and (2-8) = -6

7
-6

and sum is,

1

Upvotes: 1

Views: 153

Answers (2)

Henry Ecker
Henry Ecker

Reputation: 35686

Another approach with groupby df['no'].eq(1) (m) and m.ne(m.shift()).cumsum():

m = df['no'].eq(1)
g = (
    df.groupby([m, m.ne(m.shift()).cumsum()])['count'].first()
)
no     no
False  2      3
       4      8
True   1     10
       3      2
Name: count, dtype: int64

Then subtract the True no from the False no:

g.loc[True].values - g.loc[False].values
[ 7 -6]

and sum:

(g.loc[True].values - g.loc[False].values).sum()
1

Upvotes: 2

Mustafa Aydın
Mustafa Aydın

Reputation: 18315

Here is one way:

>>> vals = df.groupby(df.no.eq(1).diff().ne(0).cumsum())["count"].first()
>>> result = (vals * (-1) ** np.arange(len(vals))).sum()
>>> result
1

where we first groupby the consecutive groups over they are 1 or not and take the first values therein. Then to put an alternating sign to the values, we use (-1)'s power over 0..N-1. Then we sum the multiplication.

vals are this by the way:

1    10
2     3
3     2
4     8

Alternative without multiplying by -1's powers:

vals = df.groupby(df.no.eq(1).diff().ne(0).cumsum())["count"].first().to_numpy()

result = (vals[::2] - vals[1::2]).sum()

where we subtract the odd-indexed values from even-indexed values.


Another alternative based on itertools.groupby with 1 pass over data:

from itertools import groupby
from operator import itemgetter

iterable = enumerate(groupby(zip(df["no"].eq(1), df["count"]), key=itemgetter(0)))

result = sum((1-2*(j&1)) * next(gr)[1] for j, (_, gr) in iterable)

which does the sign determination on the fly with enumeration and j's parity. But not as readable.

Upvotes: 2

Related Questions