I have a dataframe with two columns:

      A B
    0 1 3
    1 2 2
    2 3 2
    3 9 3 
    4 1 1 
    ...

For a given index i, I want the rows from row i to the row j in which df.at[j,A]-df.at[i,B]>5. I don't want any rows after row j.

For example, let i=1, the output should be:

[out]
A B
2 2
3 2
9 3

Is there a simple way of do this without using loops?

CodePudding user response：

df = pd.DataFrame({'A': [10, 1, 2, 3, 9], 'B': [1, 3, 2, 2, 3]})
i = 2
base = df.at[i, 'B']
df = df.iloc[i:]
j = df[df['A'] - df.at[i, 'B'] > 5]
if not j.empty:
    print(df.iloc[:j.index[0]])
else:
    print('Condition not found')

Prints:

   A  B
2  2  2
3  3  2
4  9  3

CodePudding user response：

You could try as follows:

import pandas as pd

data = {'A': {0: 10, 1: 2, 2: 3, 3: 9}, 'B': {0: 3, 1: 2, 2: 2, 3: 3}}
df = pd.DataFrame(data)

i=1
s = df.loc[i:,'A']-df.loc[i,'B']>5
trues = s[s==True]
if not trues.empty: 
    subset = df.iloc[i:trues.idxmax() 1]
else:
    subset = pd.DataFrame()
print(subset)

   A  B
1  2  2
2  3  2
3  9  3