前提
データ加工や分析でお世話になるPythonのライブラリ「pandas」。
「forループで処理した結果を1行ずつ格納したデータフレームを作成する」
ということはよくある（と思う）のだが、なかなかスマートな書き方が分からない。
例えば以下のような感じ。
dic = {'i': [], 'i^2': [], 'i^3': [], 'i^4': []}
for i in range(10):
dic['i'].append(i)
dic['i^2'].append(i**2)
dic['i^3'].append(i**3)
dic['i^4'].append(i**4)
df = pd.DataFrame(dic)
上のコードは
| i | i^2 | i^3 | i^4 |
|---|---|---|---|
| 0 | 0 | 0 | 0 |
| 1 | 1 | 1 | 1 |
| 2 | 4 | 8 | 16 |
| 3 | 9 | 27 | 81 |
| 4 | 16 | 64 | 256 |
| 5 | 25 | 125 | 625 |
| 6 | 36 | 216 | 1296 |
| 7 | 49 | 343 | 2401 |
| 8 | 64 | 512 | 4096 |
| 9 | 81 | 729 | 6561 |
のようなデータフレームを作成するものだが、append()を4回も書くのがしんどい。
4回ぐらいなら許せるかもしれないが、これが仮に 100列とかになったら・・・
本記事では、このようなコードをいかにスマートに書くかを考えてみる。
なお、処理速度に関してはこちらの記事で議論されているのでそちらを参照されたい。
同記事に記載の通り、「作っておいたデータフレームに1行ずつ追加する」よりも「各列をリストとして作っておき、最後にデータフレーム化する」ほうが高速らしいので、本記事ではこちらの方法を採用する。
書き方いろいろ
for文でappend
keys = ['i', 'i^2', 'i^3', 'i^4']
dic = {key: [] for key in keys}
for i in range(10):
values = [i, i**2, i**3, i**4]
for key, value in zip(keys, values):
dic[key].append(value)
df = pd.DataFrame(dic)
列名を変数keys、要素をvaluesとおき、for文化してみた。
append()を書くのが1回で済み、すっきりした印象。
ただ、リストの順番をしっかりチェックしないと、意図しないデータが入れられてしまうかもしれない。
### 意図しない挙動をするパターン ###
keys = ['i', 'i^2', 'i^3', 'i^4']
dic = {key: [] for key in keys}
for i in range(10):
values = [i, i**3, i**2, i**4] # 順番間違い
for key, value in zip(keys, values):
dic[key].append(value) # i^2の列にi**3が入ってしまう
df = pd.DataFrame(dic)
こういうことを防ぐためにも、列名と追加する要素の対応が見やすいほうがよい。
ループ中に対応表を記述した上で、for文でappend
keys = ['i', 'i^2', 'i^3', 'i^4']
dic = {key: [] for key in keys}
for i in range(10):
append_data = {
'i': i,
'i^2': i**2,
'i^3': i**3,
'i^4': i**4
}
for key in keys:
dic[key].append(append_data[key])
df = pd.DataFrame(dic)
列名と追加する要素の対応が見やすい。
辞書に一括appendする関数を作る
def multi_append(dic, **kwargs):
# dic[列名].append(要素)を一括で行う
for arg in kwargs:
dic[arg].append(kwargs[arg])
dic = {key: [] for key in ['i', 'i^2', 'i^3', 'i^4']}
for i in range(10):
multi_append(dic, **{'i': i, 'i^2': i**2, 'i^3': i**3, 'i^4': i**4})
df = pd.DataFrame(dic)
上とほとんど同じことだが、appendの処理を関数multi_append()にまとめることもできる。
multi_append()は以下のように書いてもよい。
def multi_append(dic, **kwargs):
# dic[列名].append(要素)を一括で行う
for arg in dic: # argをdicのほうから取り出す
dic[arg].append(kwargs[arg])
上のように列名を辞書の方から取り出すようにしておけば、
multi_append()の引数を指定し忘れたときに例外を投げてくれる。
まとめ
記述量が少なく、かつ可読性も高いコードが書ければ最高。
ぱっと思いつくものを挙げてみたが、もっとスマートに書く方法があれば是非ご教授ください。