使用 pandas 中的 for 循環創建新列

Question

我是 python 的新手，因此也是 pandas 數據幀的新手。 假設我有以下數據集：

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

我想做的是根據我在 a 和 b 列中的信息創建新列，比如 b_1、b_2、b_3。 最終數據應如下所示：

Out[21]: 
   a  b  b_1  b_2  b_3
0  1  4    4    0    0
1  1  4    4    0    0
2  1  4    4    0    0
3  2  5    0    5    0
4  2  5    0    5    0
5  2  5    0    5    0
6  3  6    0    0    6
7  3  6    0    0    6
8  3  6    0    0    6

在 Stata 中，這是通過以下命令實現的：

forvalues i=1(1)3{
gen b_`i'=b if a==`i'
replace b_`i'=0 if b_`i'==.
}

在 python 中有類似的方法嗎？ 提前致謝

Answer 1

使用DataFrame.join與Series.unstack和DataFrame.add_prefix ：

df = df.join(df.set_index('a', append=True)['b'].unstack(fill_value=0).add_prefix('b_'))
print (df)
   a  b  b_1  b_2  b_3
0  1  4    4    0    0
1  1  4    4    0    0
2  1  4    4    0    0
3  2  5    0    5    0
4  2  5    0    5    0
5  2  5    0    5    0
6  3  6    0    0    6
7  3  6    0    0    6
8  3  6    0    0    6