具有来自相同 id 的值的 Pandas 新列，遵循条件

Question

我有一个包含多列的 DataFrame，我将向人工 df 提供代码以进行复制：

import pandas as pd
from itertools import product
df = pd.DataFrame(data=list(product([0,1,2], [0,1,2], [0,1,2])), columns=['A', 'B','C'])
df['D'] = range(len(df))

这会产生以下数据框：

   A  B  C  D
0  0  0  0  0
1  0  0  1  1
2  0  0  2  2
3  0  1  0  3
4  0  1  1  4
5  0  1  2  5
6  0  2  0  6
7  0  2  1  7
8  0  2  2  8
9  1  0  0  9

我想获得一个新列new_C，它采用 C 值，其中 B 满足条件并将其分布在 A 列中的所有匹配值上。

以下代码正是这样做的：

new_df = df[['A','B', 'D']].loc[df['C'] == 0]
new_df.columns = ['A', 'B','new_D']
df = df.merge(new_df, on=['A', 'B'], how= 'outer')

但是，我坚信有一个更好的解决方案，我不必引入一个全新的 DataFrame 并将其重新合并在一起。 最好是oneliner。

提前致谢。

期望的输出：

   A  B  C  D  new_D
0  0  0  0  0      0
1  0  0  1  1      0
2  0  0  2  2      0
3  0  1  0  3      3
4  0  1  1  4      3
5  0  1  2  5      3
6  0  2  0  6      6
7  0  2  1  7      6
8  0  2  2  8      6
9  1  0  0  9      9

编辑： 添加其他示例：

   A  B    C  D
   A  B    C  D
0  0  4  foo  0
1  0  4  bar  1
2  0  4  baz  2
3  0  5  foo  3
4  0  5  bar  4
5  0  5  baz  5
6  0  6  foo  6
7  0  6  bar  7
8  0  6  baz  8
9  1  4  foo  9

应转化为如下条件：df['C'] == 'bar'

   A  B    C  D  new_D
0  0  4  foo  0      1
1  0  4  bar  1      1
2  0  4  baz  2      1
3  0  5  foo  3      4
4  0  5  bar  4      4
5  0  5  baz  5      4
6  0  6  foo  6      7
7  0  6  bar  7      7
8  0  6  baz  8      7
9  1  4  foo  9     10

意味着所有数字都是任意的。顺序也不一样，只是碰巧取第一个数字。

Answer 1

如果你想在每次 C 等于 0 时获得一个新的基线，你可以使用：

df['new_D'] = df['D'].where(df['C'].eq(0)).ffill(downcast='infer')

旧答案

你想要什么并不完全清楚，但看起来你想重复每组 A 和 B 的第一项。你可以通过以下方式轻松实现：

df['new_D'] = df.groupby(['A', 'B'])['D'].transform('first')

如果你的数据真的由连续的整数组成，那就更简单了：

df['D'] = df['D']//3*3