具有嵌套 for 循环条件的 Python Pandas 切片

Question

我有一个这样的数据框：

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

还有一个列表，这里是 x，它的长度可能会像这样变化：

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

我想像这样根据 x 中的值对 df 进行切片

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])]

你知道一个更好的解决方案吗？像下面这样的嵌套 for 循环实际上可以工作？

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

谢谢大家！

Answer 1

试试：

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