Эффективная фильтрация в панде Dataframe

Question

У меня есть датафрейм, имеющий 3 столбца - A, B и проекты. A и B содержат идентификаторы, данные каждому ученику в классе, а проекты - это проекты, которые они делали вместе. Данные находятся на уровне столбца A и столбца B

A |  B | projects
S2   S3    5
S2   S4    2
S5   S2    1
S5   S4    1

Обратите внимание, что учащийся может появиться в любом из столбцов входного набора данных (S2 - как в коле, так и в colB) Теперь мне нужно найти количество проектов, выполненных студентом вместе со всеми другими студентами. Фрейм данных должен выглядеть следующим образом

id_ | StudentB | projects
S2     S3          5
S2     S5          1
S2     S4          2
S3     S2          5
S4     S2          2
S4     S5          1
S5     S4          1

Теперь, если я отфильтрую столбец id_ для определенного студента, все связанные идентификаторы должны быть перечислены в столбце StudentB

Мое решение ('all_student_id' - это отдельный список всех возможных идентификаторов) -

final_df = pd.DataFrame(columns = ['id_', 'studentB','projects'])
for id_ in all_student_id:
    data_ = data[(data['A']== id_) | (data['B']== id_)] 

    a = data_[['A','projects']].rename(columns= {'A':'studentB'})
    b = data_[['B','projects']].rename(columns= {'B':'studentB'})

    a_b_concat = pd.concat([a,b], axis=0)
    formatted = a_b_concat[a_b_concat['studentB']!=id_]

    temp = pd.DataFrame({'id_':[id_]*formatted.shape[0]
                        })
    temp = pd.concat([temp, formatted.reset_index(drop = True)], axis = 1)

    final_df= pd.concat([final_df, temp])

Хотя я могу сделать это, используя цикл for для всех отдельных идентификаторов студентов, я считаю, что это не масштабируемое решение, поскольку входной набор данных может быть огромным (30 миллионов строк).

Буду признателен за любую помощь в оптимизации этого решения

Answer 1

Будет ли это работать?

Итак, дайте мне знать, если этот образец работает так, как вы хотите:

m=pd.DataFrame({'A':("S2","S2","S5","S5",'S2'),'B':("S3","S4","S2","S4",'S5'), 'projects':(5,2,1,1,6)})

Это было бы примерно так:

    A   B  projects
0  S2  S3         5
1  S2  S4         2
2  S5  S2         1
3  S5  S4         1
4  S2  S5         6

Теперь я предположил, что вы хотите, чтобы, если когда-нибудь, скажем, S2 и S5 собрались вместе, либо в столбце 1, либо в столбце 2, они должны учитываться так же, как и наоборот. Воспользовавшись этой свободой, я провел некоторый анализ и получил следующее:

f=np.sort(m.iloc[:,0:2].values)
pd.concat((pd.DataFrame(f),m[['projects']]),axis=1).groupby([0,1])['projects'].sum()

Вывод, который я получил, был:

0   1 
S2  S3    5
    S4    2
    S5    7
S4  S5    1

Это переименование столбца в 0 и 1, которое можно изменить с помощью set_axis. Главное, это то, как вы хотите? Что S2 и S5, независимо от порядка, дали свою сумму в качестве выхода?

Answer 2

Другие ответы рекомендуют groupby, но из вашего желаемого результата я не согласен, что это то, что вы ищете. Кажется, вы просто хотите, чтобы отношения из B->A также были включены как отношения из A->B. Это тривиальная операция, которая может быть выполнена путем укладки в обратную сторону столбцов A и B

---

a = df.values
b = a[:, [1,0,2]].copy()

d = pd.DataFrame(np.vstack((a, b)), columns=['id_', 'StudentB', 'projects'])

  id_ StudentB projects
0  S2       S3        5
1  S2       S4        2
2  S5       S2        1
3  S5       S4        1
4  S3       S2        5
5  S4       S2        2
6  S2       S5        1
7  S4       S5        1

Теперь вы можете искать любого ученика, используя только столбец id_, хотя я бы порекомендовал pivot здесь для лучшей структуры данных:

lookp = d.pivot('id_', 'StudentB', 'projects')

StudentB   S2   S3   S4   S5
id_
S2        NaN    5    2    1
S3          5  NaN  NaN  NaN
S4          2  NaN  NaN    1
S5          1  NaN    1  NaN

Это дает вам простой способ поиска отношений между студентами, с NaN, обозначающим, что два студента не работали вместе над какими-либо проектами.

>>> lookp.loc['S2', 'S3']
5
>>> lookp.loc['S3', 'S5']
nan

Answer 3

Вы можете сделать:

# sort the students - create a new data frame
df1 = df[['A','B']]
df1.values.sort(axis=1)
df1['projects'] = df['projects']

# now groupb
df1.groupby(['A','B'])['projects'].sum().reset_index()

    A   B   projects
0   S2  S3  5
1   S2  S4  2
2   S2  S5  1
3   S4  S5  1