将pandas df转换为不同长度的列表列表

Question

I have a dataframe just like this and I need to convert it to a list of lists according to the "desired output"

    d = {'0': ['A','A','A','B','B','C','C','C','C'],
         '1': [4.3,3.2,2.1,9.1,2.0,2.8,1.7,0.8,0.2]}
    
    df = pd.DataFrame(d)
    
            0   1
        0   A   4.3
        1   A   3.0
        2   A   2.1
        3   B   9.0
        4   B   2.0
        5   C   2.8
        6   C   1.7
        7   C   0.8
        8   C   0.2
    
    # Desired output
[[4.3, 3.2, 2.1], [9.1, 2.0], [2.8, 1.7, 0.8, 0.2]]

I wrote the following to do it and it gets the job done:

d_tuples = [*list(zip(df[0],df[1]))]
keys = df[0].unique()
list_of_lists = []
for key in keys:
    list_of_lists+=[[tup[1] for tup in d_tuples if tup[0] == key]]
list_of_lists  #[[4.3, 3.2, 2.1], [9.1, 2.0], [2.8, 1.7, 0.8, 0.2]]

However, the original database is about 25,000,000 rows long and its taking some time, I was wondering if theres a more efficient way to write it.

EDIT: "desired output" means a list_of_lists where each list contains the values in column "1" for one of the unique values in column "0"

EDIT2: Added timeit results

Answer 1

groupby object is dict, you may use it to avoid agg to speed up more

In [229]: [v.tolist() for v in df.set_index('1').groupby('0').groups.values()]
Out[229]: [[4.3, 3.2, 2.1], [9.1, 2.0], [2.8, 1.7, 0.8, 0.2]]

---

Timing on 90K rows

df = pd.concat([df] * 10000)

%timeit [v.tolist() for v in df.set_index('1').groupby('0').groups.values()]
15.2 ms ± 425 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit df.groupby('0')['1'].agg(list).tolist()
32.8 ms ± 623 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [236]: %%timeit
     ...: d_tuples = [*list(zip(df['0'],df['1']))]
     ...: keys = df['0'].unique()
     ...: list_of_lists = []
     ...: for key in keys:
     ...:     list_of_lists+=[[tup[1] for tup in d_tuples if tup[0] == key]]
     ...:
69.4 ms ± 754 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Answer 2

use .groupby with .agg(list) with another .tolist() call

df = pd.DataFrame(d)
df.groupby('0')['1'].agg(list).tolist()
#out.
[[4.3, 3.2, 2.1], [9.1, 2.0], [2.8, 1.7, 0.8, 0.2]]

Answer 3

May be using apply()

list(df.groupby('0')['1'].apply(list))

#[[4.3, 3.2, 2.1], [9.1, 2.0], [2.8, 1.7, 0.8, 0.2]]

Answer 4

There probably isn't a faster way to do it with default python lists, unfortunately. Depending on data you have you might use numpy arrays (they are memory efficient, so that will give you a speed-up) — ie, list(df.groupby('0')['1'].apply(np.array)) . Depending on the number of unique keys, speed-up can be anywhere from 10% to 100% (according to local tests on my machine).

PS By the way, don't test with small dataframes. Create a bigger one like this:

N = 500
keys = np.arange(0, N)

df = {
    '0': keys[np.random.randint(0, N, size=int(1e6))],
    '1': np.random.rand(int(1e6))
}
    
df = pd.DataFrame(df)