Bin using cumulative sum rather than observations in python

Viewed 517

Let's say that I have a data frame that has a column like this:

Weight
1
1
0.75
0.5
0.25
0.5
1
1
1
1

I want to create two bins and add a column to my data frame that shows which bin each row is in, but I don't want to bin on the observations (i.e. the first 5 observations got to bin 1 and the last five to bin 2). Instead, I want to bin such that the sum of weight for each bin is equal or as close to equal as possible without changing the order of the column.

So, I want the result to be

Weight  I want  Not this
1          1       1
1          1       1
0.75       1       1
0.5        1       1
0.25       1       1
0.5        1       2
1          2       2
1          2       2
1          2       2
1          2       2

Is there something built into Pandas that already does this, or can someone share any ideas on how to make this happen? Thanks!

2 Answers

This should do it:

df = pd.DataFrame(
        {'Weight': [1, 1, 0.75, 0.5, 0.25, 0.5, 1, 1, 1, 1]})
weight_sum = df.Weight.sum()
df['bin'] = 1
df.loc[df.Weight.cumsum() > weight_sum / 2, 'bin'] = 2

print(df)

Output:

   Weight  bin
0    1.00    1
1    1.00    1
2    0.75    1
3    0.50    1
4    0.25    1
5    0.50    1
6    1.00    2
7    1.00    2
8    1.00    2
9    1.00    2

You could use pd.cut on the cumsum of the Weights column.

df = pd.DataFrame({'Weight' : [1, 1, 0.75, 0.5, 0.25, 0.5, 1, 1, 1, 1]})

s =  df['Weight'].sum()
pd.cut(df['Weight'].cumsum(), [-1, s/2, s], labels=[1,2])

For s = 8 this by default creates the groups (-1, 4] and (4, 8]. (That's mathematical notation - a value of exactly 4 would be included into the first group)

You could choose differently and put a value of exactly 4 into the second group by specifying right = False and adjusting the boundaries, which gives you the groups [0, 4) and [4, 9)

pd.cut(df['Weight'].cumsum(), [0, s/2, s+1], labels=[1,2], right=False)

The -1 and s+1 are there to specify that a value of exactly 0 or respectively 8 should still be in that group.

Related