仕事でちょくちょく必要になって、そのたびに考えたり調べたりしていたのでメモ
問題
長さ(行数)のリスト(データフレーム)を、要素数(行数)個ずつの塊に分割する。
最後の塊に分割する際に個未満になった場合は、その要素数で分割する。
リストの場合
# データの用意 lis = [1, 2, 3, 4, 5] k = 2 n = len(lis)
インデックスだけ欲しい場合
# インデックスだけ欲しい場合 [(i, i+k) for i in range(0, n, k)]
[(0, 2), (2, 4), (4, 6)]
スライスしたい場合
# リストをスライスしたい場合 [lis[i:i+k] for i in range(0, n, k)]
[[1, 2], [3, 4], [5]]
データフレームの場合
リストとはスライスのときに必要となるインデックスが変わってくる
# データの用意 import pandas as pd df = pd.DataFrame({"a": [1, 2, 3, 4, 5], "b": ["a","b","c","d","e"]}) k = 2 n = df.shape[0]
インデックスだけ欲しい場合
# インデックスだけ欲しい場合 [(i, i+k-1) for i in range(0, n, k)]
[(0, 1), (2, 3), (4, 5)]
データフレームをスライスしたい場合
# データフレームをスライスしたい場合 dfs = [df.loc[i:i+k-1, :] for i in range(0, n, k)] for df_i in dfs: display(df_i) # NOTE: display()はjupyter内でのみ有効
もうちょっとしっかり書く場合
# 表側が順番通りの整数でないデータフレームにも対応した場合 def slice_df(df: pd.DataFrame, size: int) -> list: """pandas.DataFrameを行数sizeずつにスライスしてリストに入れて返す""" previous_index = list(df.index) df = df.reset_index(drop=True) n = df.shape[0] list_indices = [(i, i+size) for i in range(0, n, size)] df_indices = [(i, i+size-1) for i in range(0, n, size)] sliced_dfs = [] for i in range(len(df_indices)): begin_i, end_i = df_indices[i][0], df_indices[i][1] begin_l, end_l = list_indices[i][0], list_indices[i][1] df_i = df.loc[begin_i:end_i, :] df_i.index = previous_index[begin_l:end_l] sliced_dfs += [df_i] return sliced_dfs
df.index = ["A","B","C","D","E"] dfs = slice_df(df, size=2) for df_i in dfs: display(df_i) # NOTE: display()はjupyter内でのみ有効