高效实现DataFrame单元格与邻域比较并构建NetworkX有向图

核心思路：用向量化运算替代双重循环 #

3000×3000的DataFrame有900万个单元格，双重循环逐个遍历、判断并添加边的方式效率极低。推荐通过numpy向量化批量生成所有符合条件的边对，再一次性导入NetworkX，从计算和IO层面大幅降低开销。

具体实现步骤 #

1. 映射坐标与单元格值 #

先把DataFrame的行/列标签转换成可批量处理的坐标元组，同时提取所有单元格的值：

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

# 假设你的DataFrame为df
# 生成所有单元格的坐标元组列表
coords = [(row, col) for row in df.index for col in df.columns]
# 提取对应的值数组，顺序与coords完全匹配
values = df.values.flatten()

2. 批量生成邻域索引 #

先把坐标转换成数字索引，方便计算邻域位置（这里以8邻域为例，可根据需求调整为4邻域）：

# 将行/列标签映射为数字索引
row_to_idx = {row: i for i, row in enumerate(df.index)}
col_to_idx = {col: j for j, col in enumerate(df.columns)}

# 把坐标转换成二维索引数组 (总单元格数, 2)
idx_array = np.array([[row_to_idx[row], col_to_idx[col]] for row, col in coords])

# 定义8邻域的偏移量（去掉对角就是4邻域）
neighbor_offsets = [(-1,-1), (-1,0), (-1,1),
                    (0,-1),          (0,1),
                    (1,-1),  (1,0), (1,1)]

3. 向量化筛选符合条件的边对 #

利用numpy广播机制批量计算邻域并判断值的大小关系：

# 扩展索引数组，为每个邻域偏移创建副本 (8, 总单元格数, 2)
neighbor_idx = idx_array[None, :, :] + np.array(neighbor_offsets)[:, None, :]

# 过滤超出DataFrame边界的无效邻域
valid_mask = (neighbor_idx[:, :, 0] >= 0) & (neighbor_idx[:, :, 0] < len(df.index)) & \
             (neighbor_idx[:, :, 1] >= 0) & (neighbor_idx[:, :, 1] < len(df.columns))

# 把有效邻域转换成values数组的一维索引
flat_neighbor_idx = neighbor_idx[:, :, 0] * len(df.columns) + neighbor_idx[:, :, 1]
flat_neighbor_idx[~valid_mask] = -1  # 标记无效邻域

# 批量比较当前值与邻域值，筛选出当前值更大的情况
current_values = values[None, :]
neighbor_values = values[flat_neighbor_idx]
edge_mask = (current_values > neighbor_values) & valid_mask

# 收集所有有效边对
edges = []
for i in range(len(neighbor_offsets)):
    src_indices = np.where(edge_mask[i])[0]
    dst_indices = flat_neighbor_idx[i][edge_mask[i]]
    edges.extend([(coords[src], coords[dst]) for src, dst in zip(src_indices, dst_indices)])

4. 批量导入NetworkX #

避免逐个调用add_edge，直接用add_edges_from一次性添加所有边：

G = nx.DiGraph()
G.add_edges_from(edges)

额外优化：稀疏矩阵降内存 #

如果符合条件的边占比不高，可通过稀疏矩阵存储邻接关系，进一步节省内存：

from scipy.sparse import coo_matrix

# 生成源和目标的一维索引
src_idx = np.repeat(np.arange(len(coords)), len(neighbor_offsets))
dst_idx = flat_neighbor_idx.flatten()

# 过滤无效边与不符合条件的边
valid_edge_mask = edge_mask.flatten() & (dst_idx != -1)
src_idx = src_idx[valid_edge_mask]
dst_idx = dst_idx[valid_edge_mask]

# 创建稀疏邻接矩阵并转换为NetworkX图
adj_matrix = coo_matrix((np.ones(len(src_idx)), (src_idx, dst_idx)), shape=(len(coords), len(coords)))
G = nx.from_scipy_sparse_array(adj_matrix, create_using=nx.DiGraph())
# 将节点索引映射回原坐标元组
nx.relabel_nodes(G, {i: coords[i] for i in range(len(coords))}, copy=False)

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