如何高效构建满足特定条件的大规模0-1对称矩阵？

首先得说，你的痛点我太懂了——15000个样本的双重循环是O(n²)量级，2亿多次Python循环，能不快才怪！而且你还把稀疏矩阵转成了密集矩阵，这不仅慢，还巨占内存，完全没必要。咱们换个思路，从分组、向量化、稀疏存储这三个点入手，能把速度提N个数量级。

问题拆解 #

你的条件其实可以拆成两个核心：

1. 必须是同一个姓氏（不同姓氏的对直接为0，不用管）
2. 同一姓氏内，两个样本的时间差小于5分钟
   而且矩阵是对称的，对角线元素全为1（因为自己和自己肯定满足条件）

优化方案：分组+向量化+稀疏矩阵 #

核心逻辑是：只在每个姓氏组内处理两两样本的关系，组间直接跳过；用numpy的向量化操作替代Python循环；用scipy的稀疏矩阵存储结果，既省内存又高效。

具体代码实现

import pandas as pd
import numpy as np
from scipy.sparse import coo_matrix, csr_matrix

date_format = '%Y-%m-%d %H:%M:%S'

# 初始化数据
df1 = pd.DataFrame([
    ['Smith', '2024-12-16 12:00:00'], 
    ['Smith', '2024-12-16 13:00:00'], 
    ['Doe', '2024-12-16 12:01:00'], 
    ['Doe', '2024-12-16 12:04:00']
])
df1.columns = ['Surname', 'Date']
# 用pd.to_datetime直接转换，比apply+lambda快得多
df1['Date'] = pd.to_datetime(df1['Date'], format=date_format)

# 存储稀疏矩阵的坐标
rows = []
cols = []

# 按姓氏分组处理，只在组内计算符合条件的对
for surname, group in df1.groupby('Surname'):
    # 获取组内样本的原始索引和时间（转成秒数，方便快速计算）
    group_idxs = group.index.values
    times_sec = group['Date'].view('int64') // 10**9  # 转成Unix时间戳（秒）
    
    # 用numpy广播计算组内所有两两时间差的绝对值
    time_diff = np.abs(times_sec[:, np.newaxis] - times_sec)
    # 筛选出时间差小于5分钟（300秒）的对，同时包含对角线（自己和自己）
    valid_pairs = np.where(time_diff < 300)
    
    # 把组内的索引映射回原始df的索引，添加到坐标列表
    rows.extend(group_idxs[valid_pairs[0]])
    cols.extend(group_idxs[valid_pairs[1]])

# 构建稀疏COO矩阵，值全为1
sparse_A = coo_matrix(
    (np.ones(len(rows), dtype=int), (rows, cols)),
    shape=(len(df1), len(df1))
)
# 转成CSR矩阵，方便后续的矩阵操作（比如乘法、切片）
sparse_A_csr = sparse_A.tocsr()

为什么这个方法快？ #

1. 减少无效计算：只在同姓氏组内处理，避免了跨姓氏的无效循环。比如如果每个姓氏平均只有10个样本，总计算量从2亿次降到15万次，直接砍了99.9%！
2. 向量化替代循环：numpy的广播和np.where都是底层C实现的，比Python的for循环快几十到上百倍。
3. 稀疏存储省资源：15000x15000的密集矩阵要占200+MB内存，而稀疏矩阵只需要存储非零元素的坐标，内存占用能降到KB级别，同时构建和操作的速度也更快。

和你原有代码的对比 #

• 你原来用lil_matrix.todense()直接生成密集矩阵，这对于15000的规模完全是灾难，内存和时间都扛不住；
• 双重Python循环的效率极低，而分组+向量化的方式把计算量压缩到了可接受的范围；
• 最终的稀疏矩阵和你原来的密集矩阵完全等价，但操作起来灵活得多。

备注：内容来源于stack exchange，提问作者JLB