盆暗の学習記録

データサイエンス ,エンジニアリング,ビジネスについて日々学んだことの備忘録としていく予定です。初心者であり独学なので内容には誤りが含まれる可能性が大いにあります。

[Python]リストやデータフレームを任意の要素数で分割する

仕事でちょくちょく必要になって、そのたびに考えたり調べたりしていたのでメモ

問題

長さ(行数) Nのリスト(データフレーム)を、要素数(行数) K個ずつの塊に分割する。

最後の塊に分割する際に K個未満になった場合は、その要素数で分割する。 f:id:nigimitama:20200125110520p:plain

リストの場合

# データの用意
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内でのみ有効

f:id:nigimitama:20200125105554p:plain

もうちょっとしっかり書く場合

# 表側が順番通りの整数でないデータフレームにも対応した場合
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内でのみ有効

f:id:nigimitama:20200125105753p:plain