python - pandas dataframe : groupby, 应用函数返回数组并映射回结果

Question

让我们考虑一个数据框:

np.random.seed(1)
df = pd.DataFrame({"x": np.random.random(size=10)})
df["y"] = np.where(df.x < 0.5, 0, 1)

输出:

          x  y
0  0.417022  0
1  0.720324  1
2  0.000114  0
3  0.302333  0
4  0.146756  0
5  0.092339  0
6  0.186260  0
7  0.345561  0
8  0.396767  0
9  0.538817  1

我想添加一个名为 z 的新列:

          x  y  z
0  0.417022  0  0
1  0.720324  1  0
2  0.000114  0  1
3  0.302333  0  2
4  0.146756  0  3
5  0.092339  0  4
6  0.186260  0  5
7  0.345561  0  6
8  0.396767  0  7
9  0.538817  1  1

我发现我可以使用 apply 函数并尝试 map 结果，但它不起作用...

z = df.groupby("y").apply(lambda d: np.arange(d["y"].size))
df["z"] = df["y"].map(z)

输出:

          x  y                         z
0  0.417022  0  [0, 1, 2, 3, 4, 5, 6, 7]
1  0.720324  1                    [0, 1]
2  0.000114  0  [0, 1, 2, 3, 4, 5, 6, 7]
3  0.302333  0  [0, 1, 2, 3, 4, 5, 6, 7]
4  0.146756  0  [0, 1, 2, 3, 4, 5, 6, 7]
5  0.092339  0  [0, 1, 2, 3, 4, 5, 6, 7]
6  0.186260  0  [0, 1, 2, 3, 4, 5, 6, 7]
7  0.345561  0  [0, 1, 2, 3, 4, 5, 6, 7]
8  0.396767  0  [0, 1, 2, 3, 4, 5, 6, 7]
9  0.538817  1                    [0, 1]

Answer 1

IIUC 你需要cumcount :

df['z'] = df.groupby('y')['y'].cumcount()
print df
          x  y  z
0  0.417022  0  0
1  0.720324  1  0
2  0.000114  0  1
3  0.302333  0  2
4  0.146756  0  3
5  0.092339  0  4
6  0.186260  0  5
7  0.345561  0  6
8  0.396767  0  7
9  0.538817  1  1

更通用的解决方案使用transform :

df["z"] = df.groupby("y")['y'].transform(lambda d: np.arange(d.size))
print df
          x  y  z
0  0.417022  0  0
1  0.720324  1  0
2  0.000114  0  1
3  0.302333  0  2
4  0.146756  0  3
5  0.092339  0  4
6  0.186260  0  5
7  0.345561  0  6
8  0.396767  0  7
9  0.538817  1  1