如何在Pandas DataFrame中高效查询祖先并计算团队平均战力

我开发了一款拥有百万级用户的游戏，每个用户默认是自身的团队队长，也可将队长更换为其他用户，因此可能出现循环依赖。现有如下结构的Pandas DataFrame：

user_id | leader_id | power
---|---|----
1 | 1 | 10
2 | 3 | 20
3 | 2 | 30
4 | 5 | 40
5 | 5 | 50

我关注的团队定义为用户自身及其层级之下的成员，示例如下：

user_id | leader_id | power | team
---|---|----|----
1 | 1 | 10 | [1]
2 | 3 | 20 | [2, 3]
3 | 2 | 30 | [2, 3]
4 | 5 | 40 | [4]
5 | 5 | 50 | [4, 5]

我需要新增一列，计算该用户团队的平均战力。但现有两种实现方法（NetworkX方案、自定义缓存方案）在5万行数据上运行速度仍不理想（分别耗时约53秒、38秒），特此寻求更高效的实现方案。测试代码如下：

import time

import networkx as nx
import numpy as np
import pandas as pd

from tests.pandas_test_utils import assert_frame_equal


def get_team_power_using_networkx(df):
    G = nx.from_pandas_edgelist(df, source='user_id', target='leader_id', create_using=nx.DiGraph())
    ancestors = {node: nx.ancestors(G, node) for node in G.nodes()}

    user_ids = df['user_id'].values

    team_ids_list = [list(ancestors.get(user_id, set()) | {user_id}) for user_id in user_ids]

    hierarchal_scores = [df.loc[df['user_id'].isin(team_ids), "power"].dropna().values for team_ids in team_ids_list]
    avg_scores = [np.sum(scores) / len(scores) if len(scores) > 0 else np.nan for scores in hierarchal_scores]
    df["team_power"] = np.where(np.isnan(avg_scores), np.nan, np.round(avg_scores).astype(int))

    return df


def get_team_power_using_custom(df):
    G = nx.from_pandas_edgelist(df, source='user_id', target='leader_id', create_using=nx.DiGraph())

    def get_ancestors(node):
        if node in ancestors:
            return ancestors[node]
        ancestors[node] = set()
        for parent in G.predecessors(node):
            ancestors[node].add(parent)
            ancestors[node].update(get_ancestors(parent))
        return ancestors[node]

    ancestors = {}
    user_ids = df['user_id'].values

    team_ids_list = [list(get_ancestors(user_id) | {user_id}) for user_id in user_ids]

    hierarchal_scores = [df.loc[df['user_id'].isin(team_ids), "power"].dropna().values for team_ids in team_ids_list]
    avg_scores = [np.sum(scores) / len(scores) if len(scores) > 0 else np.nan for scores in hierarchal_scores]
    df["team_power"] = np.where(np.isnan(avg_scores), np.nan, np.round(avg_scores).astype(int))

    return df


## Test if it works
"""
user 1 is alone. team is [1] and average power is [10]
user 2 and 3 have each other as leaders. both have teams consisting of [2, 3] and average power is [25]
user 4 reports to user 5. team is [4] and average power is [40]
user 5 reports to themself. team is [4, 5, 6, 7] and average power is [53] -> (40 + 50 + 70) / 3
user 6 reports to user 5. team is [6, 7] and average power is [70]
user 7 reports to user 6 who reports to user 5. team is [7] and average power is [70]
user 8 reports to themself. team is [8] and average power is [np.nan]
"""
user_id = [1, 2, 3, 4, 5, 6, 7, 8]
leader_id = [1, 3, 2, 5, 5, 5, 6, 8]
power = [10, 20, 30, 40, 50, np.nan, 70, np.nan]
# Create the DataFrame
data = {'user_id': user_id, 'leader_id': leader_id, 'power': power, }
df = pd.DataFrame(data)
for func in [get_team_power_using_networkx, get_team_power_using_custom]:
    print(f"Testing: {func.__name__}")
    results = func(df)[["user_id", "team_power"]]
    expected_results = pd.DataFrame({"user_id": [1, 2, 3, 4, 5, 6, 7, 8], "team_power": [10, 25, 25, 40, 53, 70, 70, np.nan]})
    assert_frame_equal(results, expected_results)

## Test speed
num_rows = 50000
np.random.seed(42)
# Generate random values for each column
user_id = np.random.randint(1, 1000, size=num_rows)
leader_id = np.random.randint(1, 1000, size=num_rows)
power = np.random.uniform(0, 100, size=num_rows)
# Create the DataFrame
data = {'user_id': user_id, 'leader_id': leader_id, 'power': power, }
df = pd.DataFrame(data)
for func in [get_team_power_using_networkx, get_team_power_using_custom]:
    start_time = time.time()  # Record the start time
    func(df)
    end_time = time.time()  # Record the end time
    print(f"Time to run {func.__name__}:", end_time - start_time)

内容的提问来源于stack exchange，提问作者elthran