面试题：Rust集合在图算法中的应用

Question

给定一个用邻接表（可以用Rust的HashMap和Vec实现）表示的无向图，实现一个函数来计算从给定起始节点到所有其他节点的最短路径（使用Dijkstra算法）。在实现过程中，需要合理使用Rust集合的特性来优化算法的空间和时间复杂度。

Answer 1

use std::collections::{BinaryHeap, HashMap};

// 定义图的边
#[derive(Eq, PartialEq)]
struct Edge {
    to: usize,
    weight: u32,
}

// 定义图
type Graph = HashMap<usize, Vec<Edge>>;

// 定义节点状态
#[derive(Debug, Eq, PartialEq)]
struct NodeState {
    node: usize,
    distance: u32,
}

// 实现从起始节点到所有其他节点的最短路径
fn dijkstra(graph: &Graph, start: usize) -> HashMap<usize, u32> {
    let mut distances = HashMap::new();
    let mut heap = BinaryHeap::new();

    // 初始化距离和堆
    for (node, _) in graph.iter() {
        if *node == start {
            distances.insert(start, 0);
            heap.push(NodeState {
                node: start,
                distance: 0,
            });
        } else {
            distances.insert(*node, u32::MAX);
        }
    }

    // Dijkstra算法核心
    while let Some(NodeState { node, distance }) = heap.pop() {
        if distance > *distances.get(&node).unwrap() {
            continue;
        }

        if let Some(edges) = graph.get(&node) {
            for edge in edges {
                let new_distance = distance + edge.weight;
                if new_distance < *distances.get(&edge.to).unwrap() {
                    distances.insert(edge.to, new_distance);
                    heap.push(NodeState {
                        node: edge.to,
                        distance: new_distance,
                    });
                }
            }
        }
    }

    distances
}

你可以使用以下方式调用这个函数：

fn main() {
    let mut graph = Graph::new();
    graph.insert(0, vec![Edge { to: 1, weight: 1 }, Edge { to: 2, weight: 4 }]);
    graph.insert(1, vec![Edge { to: 0, weight: 1 }, Edge { to: 2, weight: 2 }, Edge { to: 3, weight: 5 }]);
    graph.insert(2, vec![Edge { to: 0, weight: 4 }, Edge { to: 1, weight: 2 }, Edge { to: 3, weight: 1 }]);
    graph.insert(3, vec![Edge { to: 1, weight: 5 }, Edge { to: 2, weight: 1 }]);

    let start = 0;
    let distances = dijkstra(&graph, start);
    println!("从节点 {} 到其他节点的最短距离: {:?}", start, distances);
}