1

I would like to select rows in a dataframe based on a sum crieteria of one of the columns. For example I want the indexes of the the first rows of the dataframe where the sum of column B is less than 3:

df = pd.DataFrame({'A':[z, y, x, w], 'B':[1, 1, 1, 1]})

The only solution I have is a seperate dataframe and a while loop:

df2 = pd.DataFrame({'A':[], 'B':[]})
index = 0
while df2['B'].sum() < 3:
    df2 = df2.append(df1.loc[index])
    index += 1

The logic gets me where I need but seems unnecessarily inefficient. Does anyone have a creative way of using pandas to filter the dataframe based on sum conditional of a column?

chic9009
  • 77
  • 8

1 Answers1

2

What you describe is a cumulative sum (cumsum).

Appending rows to a DataFrame within a loop is horribly inefficient as it copies the entire DataFrame on every iteration just to append an additional small amount of data. Instead you should look to slice your original DataFrame with a Boolean mask; in this case checking where the cumsum is less than 3.

df2 = df[df['B'].cumsum().lt(3)]

#   A  B
#0  z  1
#1  y  1

df['B'].cumsum()
#0    1
#1    2
#2    3
#3    4

df['B'].cumsum().lt(3)
#0     True     <- Slicing with this Boolean Series
#1     True     <- keeps only these True rows
#2    False
#3    False
ALollz
  • 57,915
  • 7
  • 66
  • 89
  • 1
    Freaking fantastic and exactly what I was looking for. I felt the inefficiency in my chest but couldn't come up with a way to slice my way home. Thank you! – chic9009 Feb 11 '21 at 22:56
  • 1
    short follow up question. How would I add the next row. Say this represented physical goods and I wanted some buffer in my stock and in the same scenario the value for B was 2 but I wanted to include it anyway because I had a 50% margin. Here is where the while loop was nice and I don't see a one-line slicing solution to add the extra row – chic9009 Feb 11 '21 at 23:03
  • 1
    @chic9009 oh that's a good question. I won't be able to answer that right now but I'll follow up in an hour or two when I have some free time. – ALollz Feb 11 '21 at 23:11