如何切片熊猫 MultiIndex df 保留所有值直到满足特定条件？答案

【问题标题】：How to slice a pandas MultiIndex df keeping all values until a certain condition is met?如何切片熊猫 MultiIndex df 保留所有值直到满足特定条件？
【发布时间】：2019-11-15 04:25:48
【问题描述】：

我有一个 3 级 MultiIndex 数据框，我想对其进行切片，以便保留满足某个条件之前的所有值。举个例子，我有以下数据框：

                           Col1  Col2
Date          Range  Label
'2018-08-01'  1      A     900   815
                     B     850   820
                     C     800   820
                     D     950   840
              2      A     900   820
                     B     750   850
                     C     850   820
                     D     850   800

我想选择所有值，直到 Col1 变得小于 Col2。一旦我有一个 Col1

                           Col1  Col2
Date          Range  Label
'2018-08-01'  1      A     900   815
                     B     850   820
              2      A     900   820

我尝试了几个选项，但我还没有找到一个好的解决方案。我可以轻松地保留 Col1 > Col2 的所有数据：

df_new=df[df['Col1']>df['Col2']]

但这不是我需要的。我也一直在考虑循环遍历 1 级索引并使用 pd.IndexSlice 对数据帧进行切片：

idx = pd.IndexSlice
idx_lev1=df.index.get_level_values(1).unique()

for j in (idx_lev1):
    df_lev1=df.loc[idx[:,j,:],:]
    idxs=df_lev1.index.get_level_values(2)[np.where(df_lev1['Col1']<df_lev1['Col2'])[0][0]-1]
    df_sliced= df_lev1.loc[idx[:,:,:idxs],:]

然后连接各种数据帧。但是，这效率不高（我的数据框有超过 300 万个条目，所以这也是我必须考虑的问题）而且我的问题是范围索引在不同的日期重复，所以我可能不得不嵌套 2 个周期或类似的东西。

我确信一定有一个简单且更 Python 的解决方案，但我无法找到解决此问题的方法。

如果您想生成上面的数据框进行测试，您可以使用：

from io import StringIO
s="""                         
Date  Range  Label  Col1  Col2
'2018-08-01'  1  A  900   815
'2018-08-01'  1  B  850   820
'2018-08-01'  1  C  800   820
'2018-08-01'  1  D  950   840
'2018-08-01'  2  A  900   820
'2018-08-01'  2  B  750   850
'2018-08-01'  2  C  850   820
'2018-08-01'  2  D  850   800
"""
df2 = pd.read_csv(StringIO(s),
             sep='\s+',
             index_col=['Date','Range','Label'])

更新：

我尝试实现 Adam.Er8 和 Alexandre B. 的解决方案，它们与我为 SO 创建的测试数据框工作正常，但不适用于真实数据。
问题是可能存在 Col1 值总是大于 Col2 的实例，在这种情况下，我只想保留所有数据。目前提出的解决方案都不能真正解决这个问题。

对于更真实的测试用例，您可以使用以下示例：

s="""                         
Date  Range  Label  Col1  Col2
'2018-08-01'  1  1  900   815
'2018-08-01'  1  2  950   820
'2018-08-01'  1  3  900   820
'2018-08-01'  1  4  950   840
'2018-08-01'  2  1  900   820
'2018-08-01'  2  2  750   850
'2018-08-01'  2  3  850   820
'2018-08-01'  2  4  850   800
'2018-08-02'  1  1  900   815
'2018-08-02'  1  2  850   820
'2018-08-02'  1  3  800   820
'2018-08-02'  1  4  950   840
'2018-08-02'  2  1  900   820
'2018-08-02'  2  2  750   850
'2018-08-02'  2  3  850   820
'2018-08-02'  2  4  850   800
"""

或者，您可以从here 下载 hdf 文件。这是我真正使用的数据框的一个子集。

【问题讨论】：

标签： python pandas slice multi-index

【解决方案1】：

我尝试使用.cumcount() 对每一行进行编号，然后找到具有正确条件的第一行，并使用它仅过滤编号低于该值的行。

试试这个：

from collections import defaultdict

import pandas as pd
from io import StringIO

s="""
Date  Range  Label  Col1  Col2
'2018-08-01'  1  1  900   815
'2018-08-01'  1  2  950   820
'2018-08-01'  1  3  900   820
'2018-08-01'  1  4  950   840
'2018-08-01'  2  1  900   820
'2018-08-01'  2  2  750   850
'2018-08-01'  2  3  850   820
'2018-08-01'  2  4  850   800
'2018-08-02'  1  1  900   815
'2018-08-02'  1  2  850   820
'2018-08-02'  1  3  800   820
'2018-08-02'  1  4  950   840
'2018-08-02'  2  1  900   820
'2018-08-02'  2  2  750   850
'2018-08-02'  2  3  850   820
'2018-08-02'  2  4  850   800
"""
df = pd.read_csv(StringIO(s),
                 sep='\s+',
                 index_col=['Date', 'Range', 'Label'])

groupby_date_range = df.groupby(["Date", "Range"])
df["cumcount"] = groupby_date_range.cumcount()

first_col1_lt_col2 = defaultdict(lambda: len(df), df[df['Col1'] < df['Col2']].groupby(["Date", "Range"])["cumcount"].min().to_dict())

result = df[df.apply(lambda row: row["cumcount"] < first_col1_lt_col2[row.name[:2]], axis=1)].drop(columns="cumcount")
print(result)

输出：

                          Col1  Col2
Date         Range Label            
'2018-08-01' 1     1       900   815
                   2       950   820
                   3       900   820
                   4       950   840
             2     1       900   820
'2018-08-02' 1     1       900   815
                   2       850   820
             2     1       900   820

【讨论】：

顺便说一句，如果有人知道如何直接使用.min() 的结果，以后避免使用to_dict 和apply，如果你提出改进建议那就太好了:)
这些列不是索引吗？
我更新了我的问题，您的解决方案的问题是它无法处理 Col1 值总是大于 Col2 的情况（我的错是我一开始没有指定这个）。如果您能找到解决方法，我很乐意接受您的解决方案。
@baccandr 好的，完美，我编辑了我的答案，现在我的字典是一个默认字典，默认为 len(df)，所以如果现在找到 Col1

【解决方案2】：

另一种方法是使用np.where 并选择第一个索引。

groupby 中的as_index=False 让您有机会忽略groupby 中的索引列。看看这个discussion

代码：

df2 = df2.reset_index() \
         .groupby(by=["Range", "Date"], as_index=False) \
         .apply(lambda x: x.head(np.where(x.Col1 < x.Col2)[0][0])) \
         .set_index(["Date", "Range", "Label"])

print(df2)
#                           Col1  Col2
# Date         Range Label
# '2018-08-01' 1     A       900   815
#                    B       850   820
#              2     A       900   820

【讨论】：

【解决方案3】：

首先我们创建一个 "helper" 列来计算每个组。然后我们过滤我们 groupby 中的所有行 Col1 < Col2 并得到上面的 cumcount：

df2['cumcount'] = df2.groupby(level=1).cumcount()

dfs = []

for idx, d in df2.groupby(level=1):
    n = d.loc[(d['Col1'] < d['Col2']), 'cumcount'].min()-1
    dfs.append(d.loc[d['cumcount'].le(n)])

df_final = pd.concat(dfs).drop('cumcount', axis=1)

输出


                          Col1  Col2
Date         Range Label            
'2018-08-01' 1     A       900   815
                   B       850   820
             2     A       900   820

【讨论】：

【解决方案4】：

你可以这样做：

# create a dataframe with a similar structure as yours
data={
'Date': ['2019-04-08', '2019-06-27', '2019-04-05', '2019-05-01', '2019-04-09', '2019-06-19', '2019-04-25', '2019-05-18', '2019-06-10', '2019-05-19', '2019-07-01', '2019-04-07', '2019-03-31', '2019-04-01', '2019-06-09', '2019-04-17', '2019-04-27', '2019-05-27', '2019-06-29', '2019-04-24'],
'Key1': ['B', 'B', 'C', 'A', 'C', 'B', 'A', 'C', 'A', 'C', 'A', 'A', 'C', 'A', 'A', 'B', 'B', 'B', 'A', 'A'],
'Col1': [670, 860, 658, 685, 628, 826, 871, 510, 707, 775, 707, 576, 800, 556, 833, 551, 591, 492, 647, 414],
'Col2': [442, 451, 383, 201, 424, 342, 315, 548, 321, 279, 379, 246, 269, 461, 461, 371, 342, 327, 226, 467],
}

df= pd.DataFrame(data)
df.sort_values(['Date', 'Key1'], ascending=True, inplace=True)
df.set_index(['Date', 'Key1'], inplace=True)

# here the real work starts
# temporarily create a dataframe with the comparison
# which has a simple numeric index to be used later
# to slice the original dataframe
df2= (df['Col1']<df['Col2']).reset_index()

# we only want to see the rows from the first row
# to the last row before a row in which Col1<Col2
all_unwanted= (df2.loc[df2[0] == True, [0]])
if len(all_unwanted) > 0:
    # good there was such a row, so we can use it's index
    # to slice our dataframe
    show_up_to= all_unwanted.idxmin()[0]
else:
    # no, there was no such row, so just display everything
    show_up_to= len(df)
# use the row number to slice our dataframe
df.iloc[0:show_up_to]

输出是：

                 Col1  Col2
Date       Key1            
2019-03-31 C      800   269
2019-04-01 A      556   461
2019-04-05 C      658   383
2019-04-07 A      576   246
2019-04-08 B      670   442
2019-04-09 C      628   424
2019-04-17 B      551   371
--------------------------- <-- cutting off the following lines:
2019-04-24 A      414   467
2019-04-25 A      871   315
2019-04-27 B      591   342
2019-05-01 A      685   201
2019-05-18 C      510   548
2019-05-19 C      775   279
2019-05-27 B      492   327
2019-06-09 A      833   461
2019-06-10 A      707   321
2019-06-19 B      826   342
2019-06-27 B      860   451
2019-06-29 A      647   226
2019-07-01 A      707   379

【讨论】：