dleal
dleal

Reputation: 2314

cumsum pandas up to specific value - python pandas

Cumsum until value exceeds certain number:

Say that we have two Data frames A,B that look like this:

A = pd.DataFrame({"type":['a','b','c'], "value":[100, 50, 30]})
B = pd.DataFrame({"type": ['a','a','a','a','b','b','b','c','c','c','c','c'], "value": [10,50,45,10,45,10,5,6,6,8,12,10]})

The two data frames would look like this.

>>> A
  type  value
0    a    100
1    b     50
2    c     30

>>> B
   type  value
0     a     10
1     a     50
2     a     45
3     a     10
4     b     45
5     b     10
6     b      5
7     c      6
8     c      6
9     c      8
10    c     12
11    c     10

For each group in "type" in data frame A, i would like to add the column value in B up to the number specified in the column value in A. I would also like to count the number of rows in B that were added. I've been trying to use a cumsum() but I don't know exactly to to stop the sum when the value is reached,

The output should be:

  type  value
0    a      3
1    b      2
2    c      4

Thank you,

Upvotes: 1

Views: 898

Answers (2)

Divakar
Divakar

Reputation: 221764

Assuming B['type'] to be sorted as with the sample case, here's a NumPy based solution -

IDs = np.searchsorted(A['type'],B['type'])
count_cumsum = np.bincount(IDs,B['value']).cumsum()
upper_bound = A['value'] + np.append(0,count_cumsum[:-1])
Bv_cumsum = np.cumsum(B['value'])
grp_start = np.unique(IDs,return_index=True)[1]
A['output'] = np.searchsorted(Bv_cumsum,upper_bound) - grp_start + 1

Upvotes: 0

akuiper
akuiper

Reputation: 215137

Merging the two data frame before hand should help:

import pandas as pd
df = pd.merge(B, A, on = 'type')
df['cumsum'] = df.groupby('type')['value_x'].cumsum()
B[(df.groupby('type')['cumsum'].shift().fillna(0) < df['value_y'])].groupby('type').count()

# type  value
#    a      3
#    b      2
#    c      4

Upvotes: 3

Related Questions