pandasによるpivot実行後の累積和列をスマートに追加したい

###前提・実現したいこと
◆前提
・pandasで下記csvファイルを読込

【sample.csv】
id,ptn,count
001,A,1
001,B,2
001,C,3
002,A,4
002,B,5
002,C,6
003,A,7

・制約として、NaNはそのまま表示すること

◆実現したいこと
・累積和の追加自体はできたが、もっとスマートな書き方が知りたい

###発生している問題・エラーメッセージ

・ptn列の種類が増えたときのメンテナンスが面倒で、保守を考えた時に現実的なソースではない

###該当のソースコード

import pandas as pd
import numpy as np
import math

df = pd.read_csv("sample.csv", dtype=str, encoding="utf-8")
df["count"] = df["count"].astype(float)

def calc_cumsum(A, B, C):
    if math.isnan(A):
        A = 0.0
    if math.isnan(B):
        B = 0.0
    if math.isnan(C):
        C = 0.0
    return A + B + C

df.pivot(index="id", columns="ptn", values="count") \
    .reset_index() \
    .assign(cumsum = lambda df : df.apply(lambda row: calc_cumsum(row["A"], row["B"], row["C"]), axis=1))

###試したこと
上記のソースコードを実行

id	A	B	C	cumsum
001	1.0	2.0	3.0	6.0
002	4.0	5.0	6.0	15.0
003	7.0	NaN	NaN	7.0

###補足情報(言語/FW/ツール等のバージョンなど)
python 3.5.2

行動規範の内容に同意します

回答1件

ベストアンサー

pandas の pivoto_table を利用するのがよいように思います。

python
1import pandas as pd
2import numpy as np
3
4df = pd.read_csv("sample.csv", encoding="utf-8")
5pd.pivot_table(df, index=['id'], columns=['ptn'], aggfunc=np.sum, margins=True)

| id | A | B | C | All |
| -- | -- | -- | --- |
| 1 | 1.0 | 2.0 | 3.0 | 6.0 |
| 2 | 4.0 | 5.0 | 6.0 | 15.0 |
| 3 | 7.0 | NaN | NaN | 7.0 |
| All | 12.0 | 7.0 | 9.0| 28.0 |