如何并行化Python大规模数据集循环以构建邻接矩阵

Hey there! Let's fix that painfully slow adjacency matrix code and get it running efficiently—even with your huge student dataset. First, let's break down why your current code is taking so long, then walk through algorithm optimizations (the biggest win!) and parallelization for extra speed.

Why Your Current Code Is Slow #

Your triple loop (students → all courses → all courses) runs in O(N*M²) time (N=number of students, M=number of courses)—that's brutal for large datasets. On top of that, every call to have_taken_both_courses does two expensive df.loc queries, which Pandas wasn't designed to handle in tight loops. We can fix this with smarter data handling first, then add parallelism if needed.

Step 1: Optimize the Algorithm (Most Important!) #

Instead of checking every possible course pair for every student, let's group courses by students first, then only count pairs each student actually took. This cuts the complexity to O(N*K²) where K is the average number of courses per student (way smaller than M!).

Here's how to do it:

import pandas as pd
from collections import defaultdict
import itertools

# 1. Group courses by student (get a set of courses per student)
student_courses = df.groupby('Id')['Course_Title'].apply(set).to_dict()

# 2. Count all valid course pairs across students
course_pair_counts = defaultdict(int)
for courses in student_courses.values():
    # Generate all unique, unordered pairs of courses the student took
    for pair in itertools.combinations(courses, 2):
        # Sort the pair to ensure (CourseA, CourseB) and (CourseB, CourseA) are counted the same
        sorted_pair = tuple(sorted(pair))
        course_pair_counts[sorted_pair] += 1

# 3. Build the adjacency matrix from the counts
unique_classes = df['Course_Title'].unique()
vertex_dict = {course: idx for idx, course in enumerate(unique_classes)}
num_courses = len(unique_classes)

# Initialize matrix (using numpy for efficiency)
import numpy as np
adjacency_matrix = np.zeros((num_courses, num_courses), dtype=int)

# Fill the symmetric matrix
for (course1, course2), count in course_pair_counts.items():
    i = vertex_dict[course1]
    j = vertex_dict[course2]
    adjacency_matrix[i][j] = count
    adjacency_matrix[j][i] = count

This alone should make your code run orders of magnitude faster—no parallelism needed for many cases.

Step 2: Add Parallelization for Extra Speed #

If you still need more speed (e.g., 100k+ students), we can parallelize the student-level pair counting using Python's concurrent.futures.ProcessPoolExecutor. This splits the work across your CPU cores.

Here's the parallel version:

from concurrent.futures import ProcessPoolExecutor
import pandas as pd
from collections import defaultdict
import itertools
import numpy as np

def process_single_student(courses):
    """Helper function to count course pairs for one student (runs in a separate process)"""
    student_counts = defaultdict(int)
    for pair in itertools.combinations(courses, 2):
        sorted_pair = tuple(sorted(pair))
        student_counts[sorted_pair] += 1
    return student_counts

# 1. Get list of course sets per student (easier to pass to processes)
student_courses_list = list(df.groupby('Id')['Course_Title'].apply(set).values)

# 2. Parallelize the student processing
with ProcessPoolExecutor() as executor:
    # Map the helper function to each student's course set
    results = executor.map(process_single_student, student_courses_list)

# 3. Combine results from all processes
total_counts = defaultdict(int)
for student_result in results:
    for pair, count in student_result.items():
        total_counts[pair] += count

# 4. Build the adjacency matrix (same as before)
unique_classes = df['Course_Title'].unique()
vertex_dict = {course: idx for idx, course in enumerate(unique_classes)}
num_courses = len(unique_classes)
adjacency_matrix = np.zeros((num_courses, num_courses), dtype=int)

for (course1, course2), count in total_counts.items():
    i = vertex_dict[course1]
    j = vertex_dict[course2]
    adjacency_matrix[i][j] = count
    adjacency_matrix[j][i] = count

Key Tips for Parallelization #

• Always optimize the algorithm first: Parallelizing a bad algorithm still leaves you with a slow process. The initial grouping step eliminates most of the unnecessary work.
• Avoid passing large data to processes: We pass small course sets instead of the entire DataFrame to each process, which saves memory and overhead.
• Use ProcessPoolExecutor instead of ThreadPoolExecutor: Since this is CPU-bound work (counting pairs), processes are better than threads (Python's GIL limits thread performance for CPU tasks).

Bonus: Even More Optimization for Huge Datasets #

If your dataset is massive (e.g., millions of students), consider:

• Using Dask to handle out-of-core processing (work with data larger than memory).
• Converting Course_Title to a categorical type in Pandas to reduce memory usage and speed up groupby operations:

  df['Course_Title'] = pd.Categorical(df['Course_Title'])
  

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