Question

我有一个大型数据框和一个对数据框执行某些操作的函数。该函数的参数之一是掩码。

def func(df, mask):
   if mask is None:
      mask = slice(None) # first print is fast, but the second does not work.
      mask = np.ones(len(df), dtype=bool) # this makes the first print very slow. 

   print(df.loc[mask, 'A'].sum())
   high_val = (df.val > 100) & mask
   low_val = (df.val <= 100) & mask
   print(df.loc[high_val, 'A'].sum(), df.loc[low_val, 'A'].sum())

当mask为None（没有mask）时，我上面复制的代码不起作用。

当传入None时，我可以使mask成为一个所有值都等于True的numpy数组，但这会使第一次切片（df.loc [mask，'A']）的速度变慢。

我可以使mask = slice（None）。这样可以使第一个切片（df.loc [mask，'A']）快速，但随后（df.val> 100）和mask将不起作用。

我应该如何戴口罩以使其在两种情况下都能正常工作？

Answer 1

您可以重新排列逻辑表达式并规避问题：

def func(df, mask):

    if mask is None:
        print(df.A.sum())
        mask = True
    else:
        print(df.loc[mask, 'A'].sum())

    high_val = (df.val > 100) & mask
    low_val = (df.val <= 100) & mask    

    print(df.loc[high_val, 'A'].sum(), df.loc[low_val, 'A'].sum())

如何结合布尔掩码和切片对象来切片数据帧？

1 个答案: