Как группировать строки с одинаковыми значениями?

Question

У меня есть DataFrame df, как это:

     user_id  movie_id  rating
32236   1        1        5
23171   1        2        3
83307   1        3        4
62631   1        4        3
47638   1        5        3
26184   2        1        4
1333    5        1        4
172     5        2        3
54487   6        1        4
52430   7        4        5
18504   10       1        4
4617    10       4        4

Я хочу, чтобы df была сгруппирована случайным образом на основе user_id, и каждая группа с одинаковыми количествами user_id (если не может быть разделена поровну, по крайней мере, с аналогичными суммами user_id), но без объединение user_id строк.

например. делится на 2 user_id в каждой группе

     user_id  movie_id  rating
32236   1        1        5
23171   1        2        3
83307   1        3        4
62631   1        4        3
47638   1        5        3
52430   7        4        5

     user_id  movie_id  rating
26184   2        1        4
18504   10       1        4
4617    10       4        4

     user_id  movie_id  rating
1333    5        1        4
172     5        2        3
54487   6        1        4

Я написал метод group(df, n):

def group(df, n) :
    shuffled = df.sample(frac=1)
    result = np.array_split(shuffled, n) 

    dict = {}
    for i, part in enumerate(result):
        dict['df_'+str(i+1)] = part

    return dict

Но он не работает с одинаковыми номерами, которые появляются в нескольких строках, например df И я могу только указать, на сколько групп делиться, а не сколько user_id в каждой разделенной группе.

Как разделить df на группы на основе user_id, и на каждую группу с выбранными суммами user_id, как я упоминал?

Answer 1

Во-первых, получите уникальные идентификаторы пользователей и случайным образом переставьте их:

uniques = np.random.permutation(df['user_id'].unique())

Разделите уникальные идентификаторы (при условии 2 идентификатора на разделение):

splits = np.array_split(uniques, len(uniques) // 2)

Определите функцию, которая будет группировать строки на основе значения индекса:

def grouper(value):
    return np.argmax([value in split for split in splits])

И использовать его:

for group in df.set_index('user_id').groupby(grouper):
    print(group)

Answer 2

Из того, что я понимаю, вы можете использовать:

d={'df'+str(e):df[df.user_id.isin(i)] for e,i in 
    enumerate(np.array_split(np.random.permutation(df.user_id.unique()),
                                        len(df.user_id.unique())/2))}

выход

{'df0':        user_id  movie_id  rating
 26184        2         1       4
 18504       10         1       4
 4617        10         4       4, 'df1':        user_id  movie_id  rating
 32236        1         1       5
 23171        1         2       3
 83307        1         3       4
 62631        1         4       3
 47638        1         5       3
 52430        7         4       5, 'df2':        user_id  movie_id  rating
 1333         5         1       4
 172          5         2       3
 54487        6         1       4}

вы можете назвать каждый ключ в этом диктовке:

print(d['df1'])

       user_id  movie_id  rating
32236        1         1       5
23171        1         2       3
83307        1         3       4
62631        1         4       3
47638        1         5       3
52430        7         4       5