Python Pandas slice with nested for loop condition

user_13118338 0
Viewed 267

I have a dataframe like this:

df = pd.DataFrame({'A':[1,2,3], 'B':[4,5,6],'C':[7,8,9],'D':[10,11,12]})

and a list, here x, that may vary in length like this:

x = [1,4]
x = [2,5,8]

And I would like to slice df according to values in x like this

if len(x) == 1:
   df[df['A'] == x[0]]
elif len(x) == 2:
   df[(df['A'] == x[0]) & (df['B'] == x[1])]
elif len(x) == 3:
   df[(df['A'] == x[0]) & (df['B'] == x[1]) & (df['C'] == x[2])]

Do you know a better solution with nested for loop like below that actually work?

df[(df.iloc[:,i] == x[i]) for i in range(len(x))]

Thanks guys!

pandas slice python-3.x
1 Answers

Try:

def list_slice(df,x): 
    return df[df.iloc[:, :len(x)].eq(x).all(1)]

list_slice(df, [1,4])
#    A  B  C   D
# 0  1  4  7  10

list_slice(df, [2,5,8])
#    A  B  C   D
# 1  2  5  8  11
user_42384080
Related