Pandas DataFrame数组列求共同好友（交集）的性能优化方案——替代apply方法

Great question! Row-wise apply is a common culprit for slow performance in large Pandas DataFrames because it essentially runs a Python loop under the hood. Let's walk through several optimized alternatives that will speed up your common friends calculation, depending on your dataset size and future needs.

1. Fast Set-Based Intersection (Best for Medium-Sized Data) #

Set operations in Python are highly optimized, and swapping np.intersect1d with native set intersections (plus using a list comprehension instead of apply) can give you a nice performance boost.

import pandas as pd

# Your original DataFrame
test = pd.DataFrame({ 
    'person_1': ['Frodo', 'Frodo', 'Gandalf'] ,
    'person_2': ['Sam', 'Legolas', 'Legolas'] ,
    'relations_person_1': [ ['Gandalf', 'Sam', 'Legolas', 'Gollum', 'Sauron'], ['Gandalf', 'Sam', 'Legolas', 'Gollum', 'Sauron'], ['Bilbo', 'Frodo', 'Sauron', 'Sam'] ] ,
    'relations_person_2': [ ['Gandalf', 'Frodo', 'Gimli', 'Gollum'], ['Galadriel', 'Arwen', 'Gimli', 'Frodo'], ['Galadriel', 'Arwen', 'Gimli', 'Frodo'], ] 
})

# Convert relation lists to sets (avoids repeated set conversions later)
test['p1_relations_set'] = test['relations_person_1'].apply(set)
test['p2_relations_set'] = test['relations_person_2'].apply(set)

# Use list comprehension for faster row-wise processing
test['common_friends'] = [list(a & b) for a, b in zip(test['p1_relations_set'], test['p2_relations_set'])]

# Clean up temporary columns
test = test.drop(['p1_relations_set', 'p2_relations_set'], axis=1)

print(test)

Why this works: #

• List comprehensions are faster than apply because they have less overhead from Pandas' internal logic.
• Set intersections (a & b) are implemented in C, making them way faster than np.intersect1d for this use case.

2. Vectorized Explode + Group Aggregation (Best for Large/Very Large Data) #

For datasets with hundreds of thousands (or millions) of rows, we can eliminate Python loops entirely by using Pandas' vectorized operations. This approach converts list-based data into tabular form, then uses grouping to find common friends.

import pandas as pd

test = pd.DataFrame({ 
    'person_1': ['Frodo', 'Frodo', 'Gandalf'] ,
    'person_2': ['Sam', 'Legolas', 'Legolas'] ,
    'relations_person_1': [ ['Gandalf', 'Sam', 'Legolas', 'Gollum', 'Sauron'], ['Gandalf', 'Sam', 'Legolas', 'Gollum', 'Sauron'], ['Bilbo', 'Frodo', 'Sauron', 'Sam'] ] ,
    'relations_person_2': [ ['Gandalf', 'Frodo', 'Gimli', 'Gollum'], ['Galadriel', 'Arwen', 'Gimli', 'Frodo'], ['Galadriel', 'Arwen', 'Gimli', 'Frodo'], ] 
})

# Explode person_1's relations into individual rows
df_p1 = test[['person_1', 'person_2']].explode('relations_person_1').rename(columns={'relations_person_1': 'friend'})

# Explode person_2's relations into individual rows
df_p2 = test[['person_1', 'person_2']].explode('relations_person_2').rename(columns={'relations_person_2': 'friend'})

# Combine and find friends present in both lists for each (person_1, person_2) pair
combined = pd.concat([df_p1, df_p2])
common_friends = combined.groupby(['person_1', 'person_2', 'friend']).size()
common_friends = common_friends[common_friends == 2].reset_index()

# Aggregate back to lists and merge with original DataFrame
common_friends_list = common_friends.groupby(['person_1', 'person_2'])['friend'].agg(list).reset_index(name='common_friends')
test = test.merge(common_friends_list, on=['person_1', 'person_2'], how='left')

# Fill empty entries with empty lists
test['common_friends'] = test['common_friends'].fillna('[]').apply(eval)

print(test)

Why this works: #

• All operations here use Pandas' C-optimized backend, which is orders of magnitude faster than Python loops for large datasets.
• We're leveraging the power of grouping to count occurrences, which is a strength of Pandas for tabular data.

3. Graph-Based Approach (Best for Complex Relationship Analysis) #

If you plan to do more than just calculate common friends (like finding shortest paths, node centrality, or other network metrics), using a graph library like networkx makes sense.

import pandas as pd
import networkx as nx

test = pd.DataFrame({ 
    'person_1': ['Frodo', 'Frodo', 'Gandalf'] ,
    'person_2': ['Sam', 'Legolas', 'Legolas'] ,
    'relations_person_1': [ ['Gandalf', 'Sam', 'Legolas', 'Gollum', 'Sauron'], ['Gandalf', 'Sam', 'Legolas', 'Gollum', 'Sauron'], ['Bilbo', 'Frodo', 'Sauron', 'Sam'] ] ,
    'relations_person_2': [ ['Gandalf', 'Frodo', 'Gimli', 'Gollum'], ['Galadriel', 'Arwen', 'Gimli', 'Frodo'], ['Galadriel', 'Arwen', 'Gimli', 'Frodo'], ] 
})

# Build the social graph
G = nx.Graph()

# Add all nodes (people and their friends)
all_nodes = set(test['person_1']) | set(test['person_2']) | set(test['relations_person_1'].explode()) | set(test['relations_person_2'].explode())
G.add_nodes_from(all_nodes)

# Add edges for person_1's relations
for _, row in test.iterrows():
    for friend in row['relations_person_1']:
        G.add_edge(row['person_1'], friend)

# Add edges for person_2's relations
for _, row in test.iterrows():
    for friend in row['relations_person_2']:
        G.add_edge(row['person_2'], friend)

# Calculate common neighbors (common friends) for each pair
test['common_friends'] = test.apply(lambda x: list(nx.common_neighbors(G, x['person_1'], x['person_2'])), axis=1)

print(test)

Why this works: #

• NetworkX is designed for graph operations, so it's efficient for relationship-based tasks.
• This setup lets you reuse the graph for other analyses without reprocessing the data.

Performance Comparison #

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