将 DataFrame 拆分为块

Question

我有一个包含名称、年份、标签和一堆其他变量的 DataFrame。 所以它可能看起来像这样

df = pd.DataFrame({
    "name": 4*["A"] + 5*["B"],
    "year": [1999,2000,2001,2002,2010,2011,2012,2013,2014],
    "tag": [0,1,0,0,1,0,0,1,0],
    "x1": np.random.normal(size=9),
    "x2": np.random.uniform(size=9)
})

print df

  name  tag        x1        x2  year
0    A    0 -1.352707  0.932559  1999
1    A    1 -1.359828  0.724635  2000
2    A    0  1.289980  0.477135  2001
3    A    0 -0.409960  0.863443  2002
4    B    1 -1.469220  0.324349  2010
5    B    0  0.372617  0.871734  2011
6    B    0 -0.047398  0.307596  2012
7    B    1  1.240108  0.667082  2013
8    B    0  0.558432  0.284363  2014

我正在寻找一种方法将 DataFrame 分组或拆分成块，每个块都应包含

1. tag==1 和
的一行
2. tag==0、row[year+1] 和 row[year-1] 存在的所有行，row[[year+-1,"tag"]]==1 和 row[[year+-1," name"]]==row[[year,"name"]]。

简单地说，我想要大小为 3 的块，其中中间行被标记，并且被同一公司的两行未标记的行包围。 所以在上面的例子中，只有两个通过这些条件的块是

  name  tag        x1        x2  year
0    A    0 -1.352707  0.932559  1999
1    A    1 -1.359828  0.724635  2000
2    A    0  1.289980  0.477135  2001

和

7    B    0 -0.047398  0.307596  2012
8    B    1  1.240108  0.667082  2013
9    B    0  0.558432  0.284363  2014

我曾考虑按多列分组，但问题是我需要分组的行除了名称之外没有任何共同点。 我还考虑过手动引入（在 for 循环中）另一列，为每个块提供一个新 ID，然后我可以将其分组。然而，我对这种方法非常不满意，因为它看起来既不高效也不优雅。

我会很感激任何想法。

Answer 1

让我们试试这个逻辑：

df = pd.DataFrame({
    "name": 4*["A"] + 5*["B"],
    "year": [1999,2000,2001,2002,2010,2011,2012,2013,2014],
    "tag": [0,1,0,0,1,0,0,1,0],
    "x1": np.random.normal(size=9),
    "x2": np.random.uniform(size=9)
})

grp = df.groupby(['name',
                df.tag.cumsum().rolling(3, center=True, min_periods=1).max()])

chunks_df = {}
for n, g in grp:
    if g.shape[0] >= 3:
        chunks_df[n] = g

其中 chunks_df 是您分解的数据框的字典：

chunks_df[('A', 1.0)]

  name  year  tag        x1        x2
0    A  1999    0 -0.015852  0.553314
1    A  2000    1  0.367290  0.245546
2    A  2001    0  0.605592  0.524358

chunks_df[('B', 3.0)]

  name  year  tag        x1        x2
6    B  2012    0 -0.750010  0.432032
7    B  2013    1 -0.682009  0.971042
8    B  2014    0  1.066113  0.179048

详情：

• 使用 cumsum 来唯一标识/标记每个标签 == 1。
• 使用滚动窗口为 3 并获得居中的最大值 窗口，选择 -1、1 和 +1。

Answer 2

虽然@ScottBoston 的回答非常适合我在问题中给出的DataFrame，但在缺少一年的情况下它不起作用。因此，例如在

的情况下

df = pd.DataFrame({
    "name": 4*["A"] + 6*["B"],
    "year": [1999,2000,2001,2002,2008,2010,2011,2012,2013,2014],
    "tag": [0,1,0,0,0,1,0,0,1,0],
    "x1": np.random.normal(size=10),
    "x2": np.random.uniform(size=10)
})  


print df

  name  tag        x1        x2  year
0    A    0 -0.387840  0.729721  1999
1    A    1 -0.112094  0.813332  2000
2    A    0  0.913186  0.115521  2001
3    A    0 -1.088056  0.983111  2002
4    B    0  0.037521  0.743706  2008
5    B    1  0.602878  0.007256  2010
6    B    0 -0.340498  0.961602  2011
7    B    0  0.170654  0.293789  2012
8    B    1  0.973555  0.942687  2013
9    B    0 -0.643503  0.133091  2014

代码会给出

grp = df.groupby(['name',
                df.tag.cumsum().rolling(3, center=True, min_periods=1).max()])

chunks_df = {}
for n, g in grp:
    if g.shape[0] >= 3:
        chunks_df[n] = g
        print n
        print g, "\n"    


('A', 1.0)
  name  tag        x1        x2  year
0    A    0 -0.387840  0.729721  1999
1    A    1 -0.112094  0.813332  2000
2    A    0  0.913186  0.115521  2001
3    A    0 -1.088056  0.983111  2002 

('B', 2.0)
  name  tag        x1        x2  year
4    B    0  0.037521  0.743706  2008
5    B    1  0.602878  0.007256  2010
6    B    0 -0.340498  0.961602  2011 

('B', 3.0)
  name  tag        x1        x2  year
7    B    0  0.170654  0.293789  2012
8    B    1  0.973555  0.942687  2013
9    B    0 -0.643503  0.133091  2014

这表明第一个块的大小是错误的，根据原始问题中的第二个条件（年份是2008年、2010年和2011年），第二个块不应该存在。

两个人的问题是

1. 这个问题明确地保留了一行位于多个块中的可能性，因此一个额外的索引通常是不够的。
2. 必须包含年份条件，所以滚动计算需要同时在两列（标签和年份）上，根据https://stackoverflow.com/a/37491779/2336654，pandas目前不支持。

所以我现在的解决方法如下

def rolling(df, func, window_size=3):
    dxl = int(window_size/2)    
    if window_size % 2 == 0:
        dxu = dxl
    else:
        dxu = dxl+1
    xmin = dxl
    xmax = len(df)-dxu+1

    for i in xrange(xmin,xmax):
        chunk = df.iloc[i-dxl:i+dxu,:]
        if func(chunk):
            yield chunk



def valid(chunk):
    if len(chunk.name.value_counts()) != 1:
        return False
    if chunk.tag.iloc[1] != 1:
        return False
    if chunk.year.iloc[2]-chunk.year.iloc[0] != 2:
        return False
    return True



new_df = pd.DataFrame()
for ichunk, chunk in enumerate(rolling(df, window_size=3, func=valid)):
    new_df = new_df.append(chunk.assign(new_tag=ichunk), ignore_index=True)

for name, g in new_df.groupby(["name","new_tag"]):
    print name
    print g,"\n"


('A', 0)
  name  tag        x1        x2  year  new_tag
0    A    0 -1.046241  0.692206  1999        0
1    A    1  0.373060  0.919130  2000        0
2    A    0  1.316474  0.463517  2001        0 

('B', 1)
  name  tag        x1        x2  year  new_tag
3    B    0  0.376408  0.743188  2012        1
4    B    1  0.019062  0.647851  2013        1
5    B    0 -0.442368  0.506169  2014        1 

只是想我应该添加这个，以防将来有人想知道为什么接受的答案不适用于类似的问题。