python - 将 SQL 结果从自连接转换为方形 pandas 数据框

Question

这是我所追求的 tl;dr 版本；详情如下: SQL 查询为我提供了一个包含字段 [person 1 id]、[person 2 id] 和 [number of times they are in a group together] 的表。我想转换为正方形的 pandas 数据框——每人一行，每人一列，每个元素的值是它们在一个组中的次数。我正在寻找一种比遍历结果行并一次填充数据框一个元素更优雅的方法。

---

我有一个包含作业表的数据库，其中有一列用于 person_id 和一列用于 assignment_id。它还有其他内容，但对于我们的目的而言，这才是重要的:

SELECT person_id, assignment_id FROM assignments;

person_id | assignment_id
----------+--------------
      385 |           42
      163 |           29
       51 |           42
      385 |           37
      163 |           37
         ...

I want to see how often any two people have been on the same assignment. So I do:

SELECT a1.person_id AS p1_id, a2.person_id AS p2_id, COUNT(*)
FROM assignments AS a1
INNER JOIN assignments AS a2 ON a1.assignment_id = a2.assignment_id AND a1.person_id < a2.person_id
GROUP BY a1.person_id, a2.person_id

输出如下:

p1_id | p2_id | count
------+-------+------
   51 |   385 |    1
  163 |   385 |    1
         ...

现在我正在构建一个 Python 脚本来访问数据，并希望将其转储到一个数据框中，每个人一行，每个人一列，单元格包含他们共享作业的次数。所以输出会是这样的(我不关心 * 单元格中的内容 - 可以合理地为 0 或该人所做的作业数量 - 并且不关心第一行和第一列如何格式化):

 p1_id |  p_51  |  p_163 |  p_385
-------+--------+--------+--------
   51  |    *   |    0   |    1
  163  |    0   |    *   |    1
  385  |    1   |    1   |    *

我将只有大约 20 个人，因此仅一个一个地设置值不会显着影响性能，但我正在努力学习在拥有更大数据集时的良好做法。执行此类操作的正确方法是什么？

(如果这是处理它的最佳方式，我愿意修改 SQL 查询。)

Answer 1

你可以使用 groupby在将所需的列转换为 str 类型并通过加入它们并计算它们的数量进行聚合之后。

df[['person_id', 'assignment_id']] = df[['person_id', 'assignment_id']].astype(str)

df = df.groupby(['assignment_id'], as_index=False, sort=False)['person_id']   \
       .agg({'col':','.join})['col']                                         \
       .str.split(',').apply(lambda x: sorted(x, reverse=True))              \
       .apply(pd.Series).add_prefix('p_id_')                                 \
       .set_index('p_id_0', drop=False) 

您可以使用 get_dummies 进一步简化获取索引的指标变量，p_id_0，如下所示:

df1 = pd.get_dummies(df['p_id_1']).add_prefix('p_')
print (df1)

        p_163  p_385
p_id_0              
51        0.0    1.0
163       0.0    0.0
385       1.0    0.0

df2 = pd.get_dummies(df['p_id_0']).add_prefix('p_')
print (df2)

        p_163  p_385  p_51
p_id_0                    
51        0.0    0.0   1.0
163       1.0    0.0   0.0
385       0.0    1.0   0.0

然后，在将索引帧的所有值映射到 0 之后连接这些单独的 dataframes，然后将名为 columns 的同一列分组在一起:

df_final = pd.concat([df1, df2.applymap(lambda x: 0)], axis=1).add_prefix('p_')
print (df_final.groupby(df.columns, axis=1).sum())

        p_163  p_385  p_51
p_id_0                    
51        0.0    1.0   0.0
163       0.0    0.0   0.0
385       1.0    0.0   0.0