Numba JIT编译中替代变长列表字典的查找数据结构咨询

问题描述 #

需要对处理网络图的Python函数使用Numba即时编译，但代码中用于存储节点邻接点的字典（键为int类型，值为长度各异的int列表）无法被Numba支持。现有两种思路存在缺陷：

• 每次从配对列表提取邻接点会大幅降低运行速度，违背使用Numba的初衷；
• 使用二维numpy数组会因行长度不一致导致内存浪费严重。

可行解决方案 #

方案1：使用Numba Typed容器替代原生Python容器 #

Numba提供了typed.Dict和typed.List专门用于JIT编译场景，支持变长列表作为字典值，只需提前指定类型即可。

修改后的可编译代码：

from numba.typed import Dict, List
from numba import types, jit

@jit(nopython=True)
def f():
    pairs = [(0, 1), (2, 0)]
    # 初始化Numba类型的字典：键为int64，值为int64类型的变长列表
    d = Dict.empty(
        key_type=types.int64,
        value_type=List.empty_list(types.int64)
    )

    def add_element_to_dict_list(key, element, dict_list):
        if key in dict_list:
            dict_list[key].append(element)
        else:
            new_list = List.empty_list(types.int64)
            new_list.append(element)
            dict_list[key] = new_list

    for p1, p2 in pairs:
        add_element_to_dict_list(p1, p2, d)
        add_element_to_dict_list(p2, p1, d)

    # 遍历输出字典内容
    for key in d:
        print(f"{key}: {list(d[key])}")

f()

该方案与原代码逻辑几乎一致，无需大幅重构，同时能充分利用Numba的JIT加速。

方案2：压缩邻接表存储（偏移量数组+值数组） #

这是图计算中常用的高效存储方式，通过两个numpy数组实现：

• offsets数组：offsets[i]表示节点i的邻接点在values数组中的起始索引；
• values数组：按顺序存储所有节点的邻接点。

代码示例：

import numpy as np
from numba import jit

@jit(nopython=True)
def build_compressed_adjacency(pairs, num_nodes):
    # 统计每个节点的邻接数量
    counts = np.zeros(num_nodes, dtype=np.int64)
    for p1, p2 in pairs:
        counts[p1] += 1
        counts[p2] += 1
    
    # 计算偏移量
    offsets = np.zeros(num_nodes + 1, dtype=np.int64)
    for i in range(num_nodes):
        offsets[i+1] = offsets[i] + counts[i]
    
    # 填充邻接值
    values = np.zeros(offsets[-1], dtype=np.int64)
    current = np.zeros(num_nodes, dtype=np.int64)
    for p1, p2 in pairs:
        pos = offsets[p1] + current[p1]
        values[pos] = p2
        current[p1] += 1
        
        pos = offsets[p2] + current[p2]
        values[pos] = p1
        current[p2] += 1
    
    return offsets, values

# 使用示例
pairs = [(0,1), (2,0)]
num_nodes = 3
offsets, values = build_compressed_adjacency(pairs, num_nodes)

# 提取节点0的邻接点
print(f"节点0的邻接点: {values[offsets[0]:offsets[1]]}")
# 提取节点1的邻接点
print(f"节点1的邻接点: {values[offsets[1]:offsets[2]]}")

该结构内存利用率高，Numba对numpy数组的支持极佳，访问速度快，适合节点总数已知的场景，是性能最优的选择。

方案3：预分配固定大小二维数组（适合最大邻接数可预估场景） #

若能提前确定所有节点中最多的邻接点数量，可创建(num_nodes, max_neighbors)的二维数组，用特殊值（如-1）标记空位置。此方案存在一定内存浪费，但实现简单，适合最大邻接数较小的场景。

示例代码片段：

import numpy as np
from numba import jit

@jit(nopython=True)
def build_fixed_adjacency(pairs, num_nodes, max_neighbors):
    adj = np.full((num_nodes, max_neighbors), -1, dtype=np.int64)
    counts = np.zeros(num_nodes, dtype=np.int64)
    
    for p1, p2 in pairs:
        adj[p1, counts[p1]] = p2
        counts[p1] += 1
        adj[p2, counts[p2]] = p1
        counts[p2] += 1
    
    return adj

# 使用示例
max_neighbors = 2
adj = build_fixed_adjacency(pairs, num_nodes, max_neighbors)
print(f"节点0的邻接点: {adj[0][adj[0] != -1]}")

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