Pandas：合并两个数据框时保持主数据框的索引不变

Question

我正在尝试合并两个数据框：df 和 relations_df

df: (df 的索引是帖子的 ID)

df 的示例：

Post     similar_posts
-----     --------
Head-1    [123046, 1915, 111357]
Head-2    [654, 239]

relations_df 的示例：

id   related_id   keywords
---   ---------   --------
17    123046      ['key1', 'key2']
17    1915        ['key1', 'key2', 'key6']
17    111357      ['key3', 'key4']
18    654         ['key4', 'key6', 'key1']
18    239         ['key5', 'key2']

期望的输出：(df 的索引应该保持不变)

Title     similar_posts           keywords
-----     --------                --------
Head-1    [123046, 1915, 111357]  [['key1', 'key2'], ['key1', 'key2', 'key6'], ['key3', 'key4']]
Head-2    [654, 239]              [['key4', 'key6', 'key1'], ['key5', 'key2']]

我对 pandas 还不是很了解，所以我试着问 ChatGPT，但没有得到满意的答案...

这是我目前的代码：

def addkeywords(args):
    df, relations_df = args
    
    # Merge relations_df with df on 'id' and 'related_id'
    merged_df = pd.merge(df, relations_df, left_on='id', right_on='related_id', suffixes=('_left', '_right'), how='left')
    
    # Group by 'id' and aggregate 'keywords' into lists
    grouped_df = merged_df.groupby('id', as_index=False)['keywords'].agg(list)
    # grouped_df.rename(columns={'keywords_right': 'keywords'}, inplace=True)
    
    # Merge back to the original DataFrame
    df = pd.merge(df, grouped_df, on='id', how='left')
    
    # Fill NaN values in 'keywords' column with empty lists
    df['keywords'].fillna(value='', inplace=True)
    
    return df

'df' 数据框的索引应该保持一致（因为索引代表了 df 中的帖子 ID）

'relations_df' 的索引是自动生成的，所以可能会变化。

有没有人能帮我一下呢？

Answer 1

Pandas这个工具并不是专门用来处理像列表这样的嵌套结构的。虽然在纯Pandas中可以通过一些操作（比如拆分、合并、再聚合）来实现，但这样做效率不高。最简单的方法还是使用字典，然后对字典的键进行循环。

d = (relations_df
     .set_index('related_id')['keywords']
     .to_dict()
     )

df['keywords'] = [[d.get(x) for x in l]
                  for l in df['similar_posts']]

需要注意的是，所有创建的子列表都是和字典中的那些共享的。如果你打算修改它们，最好先做一个副本：

df['keywords'] = [[d.get(x).copy() for x in l]
                  for l in df['similar_posts']]

输出：

     Post           similar_posts                                          keywords
0  Head-1  [123046, 1915, 111357]  [[key1, key2], [key1, key2, key6], [key3, key4]]
1  Head-2              [654, 239]                [[key4, key6, key1], [key5, key2]]

Answer 2

代码

有一种可能的方法，虽然可能还有更好的办法，就是用 explode 和 merge 把数据合并成一个列表。

out = (df
       .explode('similar_posts')
       .merge((relations_df.explode('keywords')
               .rename({'related_id': 'similar_posts'}, axis=1)
               .drop('id', axis=1)
              ), how='left')
       .groupby('similar_posts', as_index=False, sort=False)
       .agg({'Post': 'first', 'keywords': list})
       .groupby('Post', as_index=False, sort=False).agg(list)
      )

输出结果

---

示例代码

import pandas as pd

data1 = {'Post': ['Head-1', 'Head-2'], 'similar_posts': [[123046, 1915, 111357], [654, 239]]}
data2 = {'id': [17, 17, 17, 18, 18], 'related_id': [123046, 1915, 111357, 654, 239], 'keywords': [['key1', 'key2'], ['key1', 'key2', 'key6'], ['key3', 'key4'], ['key4', 'key6', 'key1'], ['key5', 'key2']]}

df = pd.DataFrame(data1)
relations_df = pd.DataFrame(data2)

Answer 3

# Group by 'id' and aggregate 'related_id' and 'keywords'
df_grouped = relations_df.groupby('id').agg({'related_id': list, 'keywords': list}).reset_index()
df_grouped.set_index('id', inplace=True)
  
# Merge on 'id' of relations_df and 'index' of df
merged_df = pd.merge(df, df_grouped, left_index=True, right_index=True)

我用上面的代码成功得到了想要的结果。

1. 把扁平化后的 relations_df 压缩了一下，并把它的索引设置为 'id'
2. 根据索引把 'df' 和 'relations_df' 合并在一起