每個 pandas dataframe 切片的元組列表

Question

我需要做一些與這個問題非常相似的事情： Pandas convert dataframe to array of tuples

不同之處在於，我不僅需要獲取整個 DataFrame 的單個元組列表，還需要獲取根據某些列值切片的元組列表列表。

假設這是我的數據集：

         t_id  A    B
         ----- ---- -----
    0    AAAA     1   2.0
    1    AAAA     3   4.0
    2    AAAA     5   6.0
    3    BBBB     7   8.0
    4    BBBB     9  10.0
    ...

我想生成 output：

        [[(1,2.0), (3,4.0), (5,6.0)],[(7,8.0), (9,10.0)]]

也就是說，一個列表為“AAAA”，另一個列表為“BBBB”，依此類推。

我試過兩個嵌套的 for 循環。 它似乎有效，但花費的時間太長（實際數據集有 ~1M 行）：

    result = []
    for t in df['t_id'].unique():
        tuple_list= []
        
        for x in df[df['t_id' == t]].iterrows():
            row = x[1][['A', 'B']]
            tuple_list.append(tuple(x))
        
        result.append(tuple_list)

有更快的方法嗎？

Answer 1

您可以按列t_id ，遍歷組並將每個子 dataframe 轉換為元組列表：

[g[['A', 'B']].to_records(index=False).tolist() for _, g in df.groupby('t_id')]
# [[(1, 2.0), (3, 4.0), (5, 6.0)], [(7, 8.0), (9, 10.0)]]

Answer 2

我認為這也應該有效：

import pandas as pd
import itertools


df = pd.DataFrame({"A": [1, 2, 3, 1], "B": [2, 2, 2, 2], "C": ["A", "B", "C", "B"]})

tuples_in_df = sorted(tuple(df.to_records(index=False)), key=lambda x: x[0])
output = [[tuple(x)[1:] for x in group] for _, group in itertools.groupby(tuples_in_df, lambda x: x[0])]
print(output)

出去：

[[(2, 'A'), (2, 'B')], [(2, 'B')], [(2, 'C')]]