Создание нового столбца в кадре данных pandas со списком значений из другого столбца без использования «groupby»

Question

Я работаю с большими наборами данных, делая функции панды и групповые функции занимают много времени/используют слишком много памяти. Я слышал, что некоторые люди говорят, что groupby может быть медленным, но мне трудно найти лучшее решение.

Если мой dataframe имеет 2 колонки, похожие на:

df = pd.DataFrame({'a':[1,2,2,4], 'b':[1,1,1,1]}) 

a  b  
1  1  
2  1  
2  1  
4  1  

Я хочу, чтобы возвращать список значений, которые соответствуют к значению в другой колонке:

a  b  list_of_b 
1  1  [1] 
2  1  [1,1] 
2  1  [1,1] 
4  1  [1] 

настоящее время я использую:

df_group = df.groupby('a') 
df['list_of_b'] = df.apply(lambda row: df_group.get_group(row['a'])['b'].tolist(), axis=1) 

Код выше работает для мелочей, но не на больших кадрах данных (df> 1,000,000 строк) У вас более быстрый способ сделать это?

Answer 1

Кратчайший решение, которое я могу придумать:

df = pd.DataFrame({'a':[1,2,2,4], 'b':[1,1,1,1]}) 
df.join(pd.Series(df.groupby(by='a').apply(lambda x: list(x.b)), name="list_of_b"), on='a') 

    a b list_of_b 
0 1 1  [1] 
1 2 1 [1, 1] 
2 2 1 [1, 1] 
3 4 1  [1] 

Answer 2

На ряду 4K ДФ я получаю следующее:

In [29]: 
df_group = df.groupby('a') 
​ 
%timeit df.apply(lambda row: df_group.get_group(row['a'])['b'].tolist(), axis=1) 
%timeit df['a'].map(df.groupby('a')['b'].apply(list)) 

1 loops, best of 3: 4.37 s per loop 
100 loops, best of 3: 4.21 ms per loop 

Answer 3

Просто делать группировку, а затем присоединиться обратно к первоначальному dataframe, кажется, совсем немного быстрее:

def make_lists(df): 
    g = df.groupby('a') 
    def list_of_b(x): 
     return x.b.tolist() 
    return df.set_index('a').join(
     pd.DataFrame(g.apply(list_of_b), 
        columns=['list_of_b']), 
     rsuffix='_').reset_index() 

Это дает я 192ms на петле с 1M строк генерируется так:

df1 = pd.DataFrame({'a':[1,2,2,4], 'b':[1,1,1,1]}) 
low = 1 
high = 10 
size = 1000000 
df2 = pd.DataFrame({'a':np.random.randint(low,high,size), 
        'b':np.random.randint(low,high,size)}) 

make_lists(df1) 
Out[155]: 
    a b list_of_b 
0 1 1 [1] 
1 2 1 [1, 1] 
2 2 1 [1, 1] 
3 4 1 [1] 
In [156]: 


%%timeit 
make_lists(df2) 
10 loops, best of 3: 192 ms per loop