使用冒号(:)运算符通过函数调用创建DataFrame切片

时间:2018-03-30 12:02:28

标签: python pandas dataframe

我想定义一个将索引返回到DataFrame的函数。例如,我有

import pandas as pd

df = pd.DataFrame({'A':['a', 'a', 'b'], 'B':[1,2,3], 'C':[1.1, 2.2 ,3.3]})
df = df.set_index(['A', 'B'])

然后,我可以通过df.loc['a',1]df.loc['a',:]切入DataFrame,这将返回C中与索引AB上的条件匹配的值。在第一种情况下,我会得到一个值,在第二种情况下是两个。

我不想每次都明确地编写索引,而是想定义一个自动生成它的函数。我的伪函数(不起作用)可能是,

def func(df, a, b=None):
    if b is None:
        b = ':'
    return df.loc[a,b]

有了这个,我希望能够像func(df, 'a', 1)func(df, 'a', None)这样做,就像分别直接调用df.loc['a',1]df.loc['a',:]一样。

由于我的索引中有多个级别,如果传递给函数的参数为​​:,则有一种定义None运算符的方法非常有用。

更新 我在索引中有很多级别,即我希望避免为if的每个可能组合写一个None语句。例如,假设我在索引中有四个级别。做一些像这样的事情会很棒

# Does not work
def func(df, a, b=None, c=None, d=None):
    if b is None:
        b = ':'
    if c is None:
        c = ':'
    if d is None:
        d = ':'
    return df.loc[a,b,c,d]

而不是为参数中的None的每个组合写一个if语句(在本例中为7)。

1 个答案:

答案 0 :(得分:1)

问题是您使用当前函数将':'作为字符串传递给df.loc。这样的事情应该解决问题:

def func(df, a, b=None):
    if b is None:
        b = slice(None)
    return df.loc[a, b]

我的理解是,在这里使用slice(None)相当于传递:。我只测试了你的虚拟数据集,但是值得对你的真实数据进行测试。