pandas 计算互为好友的用户数

看过来

《pandas 教程》持续更新中，提供建议、纠错、催更等加作者微信: gairuo123（备注：pandas教程）和关注公众号「盖若」ID: gairuo。跟作者学习，请进入 Python学习课程。欢迎关注作者出版的书籍：《深入浅出Pandas》和《Python之光》。

在社交业务中，我们经常要分析用户之间的关系情况，其中互为好友是一个关注方向。在本例中，我们使用 Python 对原始数据进行处理，以获得互为好友的对数。

数据和需求

数据源如下，id 为用户，friend_id 为此用户的好友列表：

import pandas as pd
from io import StringIO

data = '''
id  friend_id
111 222,333,444
222 555
333 666,777,999
444 222,999,111
555 111,666
666 333
123 3434,111
999 211,333,133
133 111,999
'''

df = pd.read_csv(StringIO(data), sep=r'\s+')
df
'''
    id    friend_id
0  111  222,333,444
1  222          555
2  333  666,777,999
3  444  222,999,111
4  555      111,666
5  666          333
6  123     3434,111
7  999  211,333,133
8  133      111,999
'''

现在需要根据以上用户的好友列表情况计算得到数据中有多少对互为好友的用户。

思路分析

首先，源数据不是一个规范的关系型数据结构，它将用户的好友对应到一个列表中，虽然 id 列是唯一的，但好友列不是纯用户id形式，需要我们处理成一个规范的用户关系表，再通过这个关系表做下一步计算。

在用户关系表的基础上，我们先将用户列和好友列连接了一个新列，新列用 - 连接，表示两个用户有好友关系，在连接时要并指定先后，如升序，得到这样的格式： 111-222、222-999 等，上小的数字在前边，保证有相同好友关系。

然后取重复值数量，大于 1 就说明是互相关注。

代码

首先将数据转为关系型结构，可以使用 df.explode() 方法，它能将类似列表结构的数据中的每个元素转换为独立的一行。在用 explode 转换之前可以先将 friend_id 转换为列表形式：

df.assign(friend_id=lambda x: x.friend_id.str.split(','))
'''
    id        friend_id
0  111  [222, 333, 444]
1  222            [555]
2  333  [666, 777, 999]
3  444  [222, 999, 111]
4  555       [111, 666]
5  666            [333]
6  123      [3434, 111]
7  999  [211, 333, 133]
8  133       [111, 999]
'''

使用 explode：

(
    df.assign(friend_id=lambda x: x.friend_id.str.split(','))
    .explode('friend_id')
)
'''
    id friend_id
0  111       222
0  111       333
0  111       444
1  222       555
2  333       666
2  333       777
2  333       999
3  444       222
3  444       999
3  444       111
4  555       111
4  555       666
5  666       333
6  123      3434
6  123       111
7  999       211
7  999       333
7  999       133
8  133       111
8  133       999
'''

接思路转换为关系列：

(
    df.assign(friend_id=lambda x: x.friend_id.str.split(','))
    .explode('friend_id')
    .agg(lambda x: pd.Series(list(x), dtype=str).sort_values().to_list(), axis=1)
)
'''
0     [111, 222]
0     [111, 333]
0     [111, 444]
1     [222, 555]
2     [333, 666]
2     [333, 777]
2     [333, 999]
3     [222, 444]
3     [444, 999]
3     [111, 444]
4     [111, 555]
4     [555, 666]
5     [333, 666]
6    [123, 3434]
6     [111, 123]
7     [211, 999]
7     [333, 999]
7     [133, 999]
8     [111, 133]
8     [133, 999]
dtype: object
'''

计算每个关系的重复值：

(
    df.assign(friend_id=lambda x: x.friend_id.str.split(','))
    .explode('friend_id')
    .agg(lambda x: pd.Series(list(x), dtype=str).sort_values().to_list(), axis=1)
    .astype(str)
    .value_counts()
)
'''
['133', '999']     2
['111', '444']     2
['333', '999']     2
['333', '666']     2
['111', '555']     1
['222', '555']     1
['111', '333']     1
['111', '222']     1
['444', '999']     1
['111', '133']     1
['111', '123']     1
['333', '777']     1
['555', '666']     1
['222', '444']     1
['123', '3434']    1
['211', '999']     1
dtype: int64
'''

最后筛选出数量大于 1 的并计数：

(
    df.assign(friend_id=lambda x: x.friend_id.str.split(',')) # 转为列表
    .explode('friend_id') # 转换为独立行
    # 拼接关系字符串
    .agg(lambda x: pd.Series(list(x), dtype=str).sort_values().to_list(), axis=1)
    .astype(str)
    .value_counts() # 计算不重复值
    .loc[lambda x: x>1] # 筛选关系大于 1 的
    .pipe(len) # 统计数量
)
# 4

最后可得互相好友的共有 4 对用户。

（完）

pandas 计算互为好友的用户数

数据和需求

思路分析

代码

相关内容