python - 在特定条件下聚合 Dataframe 中的行值

Question

我有一个包含两列的 dataframe。我想构建所有连续行的总和，其中 column 0 的值为 -1。

我的 dataframe 看起来……像这样:

 0   2
 1   3
-1   4
-1   7
 0   2
-1   0
-1   1
-1   3
 5   0

所需的输出应该是这样的:

 0   2
 1   3
-1   11
 0   2
-1   4
 5   0

第二列中的值都等于或大于零。如果有帮助，第一列的值等于或大于 -1。我的方法是一个循环，我在其中创建第二个 dataframe，然后推回不等于 -1 的每个值并累加，当我找到一个 -1，但我想，这种方法效率不高。伪代码:

sum = 0
found = False
for row in dataframe:
   if row[0] != -1:
       if found:
           new_df.append([-1, sum])
           sum = 0
           found = False
       new_df.append(row)
   elif row[0] == -1:
       found = True
       sum += row[1]

是否有内置的 python 或 pandas 函数可以用来实现我的目标？

Answer 1

在我看来，这里有必要创建 Series 以 100% 确定 helper 组的值是否与 -1 之外的值相同，因此添加了 0.5 到 index 以区分:

df = df.reset_index(drop=True)

m = df['a'] == -1
s = m.ne(m.shift()).cumsum()[m].reindex(df.index).fillna(df.index.to_series().add(.5))
df = df.groupby(s).agg({'a':'first', 'b':'sum'}).reset_index(drop=True)
print (df)
   a   b
0  0   2
1  1   3
2 -1  11
3 -1   4
4  0   2
5  5   0

解释:

df = pd.DataFrame({'a': [0, 1, -1, -1, 0, -1, -1, -1, 5],
                   'b': [2, 3, 4, 7, 2, 0, 1, 3, 0]})
print (df)
   a  b
0  0  2
1  1  3
2 -1  4
3 -1  7
4  0  2
5 -1  0
6 -1  1
7 -1  3
8  5  0

如有必要，首先创建默认索引，因为解决方案中使用了唯一索引值:

df = df.reset_index(drop=True)

然后为 -1 和其他值创建连续的组:

m = df['a'] == -1
print (df.assign(groups = m.ne(m.shift()).cumsum()))
   a  b  groups
0  0  2       1
1  1  3       1
2 -1  4       2
3 -1  7       2
4  0  2       3
5 -1  0       4
6 -1  1       4
7 -1  3       4
8  5  0       5

然后使用 boolean indexing 仅过滤 -1 值(通过掩码 b)和另一个不匹配的值通过 Series.reindex 转换为 NaNs :

m = df['a'] == -1
print (df.assign(groups = m.ne(m.shift()).cumsum(),
                 filtered = m.ne(m.shift()).cumsum()[m].reindex(df.index)))

   a  b  groups  filtered
0  0  2       1       NaN
1  1  3       1       NaN
2 -1  4       2       2.0
3 -1  7       2       2.0
4  0  2       3       NaN
5 -1  0       4       4.0
6 -1  1       4       4.0
7 -1  3       4       4.0
8  5  0       5       NaN

然后用 0.5 的索引值替换缺失值 - -1 的组和替换的 NaN 的值永远不会发生冲突:

m = df['a'] == -1
print (df.assign(groups = m.ne(m.shift()).cumsum(),
                 filtered = m.ne(m.shift()).cumsum()[m].reindex(df.index),
                 idx = df.index.to_series().add(.5),
                 groups1 = m.ne(m.shift()).cumsum()[m].reindex(df.index).fillna(df.index.to_series().add(.5))))

   a  b  groups  filtered  idx  groups1
0  0  2       1       NaN  0.5      0.5
1  1  3       1       NaN  1.5      1.5
2 -1  4       2       2.0  2.5      2.0
3 -1  7       2       2.0  3.5      2.0
4  0  2       3       NaN  4.5      4.5
5 -1  0       4       4.0  5.5      4.0
6 -1  1       4       4.0  6.5      4.0
7 -1  3       4       4.0  7.5      4.0
8  5  0       5       NaN  8.5      8.5

然后是帮助程序 Series 传递给 groupby 并聚合 sum 用于第二列和 first 用于第一列GroupBy.agg最后一个由 DataFrame.reset_index 删除索引使用 drop=True:

df = df.groupby(s).agg({'a':'first', 'b':'sum'}).reset_index(drop=True)
print (df)
   a   b
0  0   2
1  1   3
2 -1  11
3 -1   4
4  0   2
5  5   0

---

另一种解决方案，更简单且性能更好:

df = df.reset_index(drop=True)

m = df['a'] == -1
s = df.reset_index()
      .groupby(m.ne(m.shift()).cumsum()[m])
      .agg({'index':'first', 'b':'sum'})
      .set_index('index')
      .assign(a = -1)

df = df[~m].append(s, sort=True).sort_index()
print (df)
   a   b
0  0   2
1  1   3
2 -1  11
4  0   2
5 -1   4
8  5   0

解释:

首先是必要的默认索引:

df = df.reset_index(drop=True)

然后将 -1 的列与 bool 掩码进行比较:

m = df['a'] == -1

通过 reset_index 将索引转换为列，不使用 drop 参数:

print (df.reset_index())
   index  a  b
0      0  0  2
1      1  1  3
2      2 -1  4
3      3 -1  7
4      4  0  2
5      5 -1  0
6      6 -1  1
7      7 -1  3
8      8  5  0

使用 shift 和 cumsum 创建连续的组，并按 -1 组的掩码过滤:

print (m.ne(m.shift()).cumsum()[m])
2    2
3    2
5    4
6    4
7    4
Name: a, dtype: int32

按索引列聚合first，按b列聚合sum:

print (df.reset_index()
        .groupby(m.ne(m.shift()).cumsum()[m])
        .agg({'index':'first', 'b':'sum'}))

     index   b
a             
2.0      2  11
4.0      5   4

将 index 列转换为 index by DataFrame.set_index :

print(df.reset_index()
      .groupby(m.ne(m.shift()).cumsum()[m])
      .agg({'index':'first', 'b':'sum'})
      .set_index('index'))
        b
index    
2      11
5       4

通过 DataFrame.assign 添加 a 列，常量为 -1 :

s = (df.reset_index()
      .groupby(m.ne(m.shift()).cumsum()[m])
      .agg({'index':'first', 'b':'sum'})
      .set_index('index')
      .assign(a = -1))
print (s)
        b  a
index       
2      11 -1
5       4 -1

最后按boolean indexing过滤掉-1行通过 ~ 的反向掩码:

print (df[~m])
  a  b
0  0  2
1  1  3
4  0  2
8  5  0

然后通过DataFrame.append向原始数据添加新数据:

print (df[~m].append(s, sort=True))
   a   b
0  0   2
1  1   3
4  0   2
8  5   0
2 -1  11
5 -1   4

最后DataFrame.sort_index对于相同的顺序:

print (df[~m].append(s, sort=True).sort_index())
   a   b
0  0   2
1  1   3
2 -1  11
4  0   2
5 -1   4
8  5   0