Skip to content

Instantly share code, notes, and snippets.

@wwylele
Last active July 24, 2020 04:02
Show Gist options
  • Save wwylele/2405946678b42dc7236f85f640bbdfad to your computer and use it in GitHub Desktop.
Save wwylele/2405946678b42dc7236f85f640bbdfad to your computer and use it in GitHub Desktop.
dijkstra template in rust
use std::ops::{Index, IndexMut, Add};
use std::cmp::Ord;
#[derive(Clone, Copy)]
pub enum NodeStatus {
Unknown,
Visiting(usize),
Visited,
}
impl Default for NodeStatus {
fn default() -> NodeStatus {
NodeStatus::Unknown
}
}
/// Dijkstra algorithm.
///
/// Node: type for graph node ID.
/// Edge: type for graph edge, typically containing edge weight.
/// Path: type for a path in the graph, typically containing the path weight.
/// - implements Ord, and must have total order.
/// - implements &Path + Edge -> Path, and must be non-decreasing on Path.
/// Board: an efficitient container mapping Node to NodeStatus.
/// - implements Board[&Node] -> &mut NodeStatus, and must be deterministic.
/// start_node: the starting node for all paths.
/// start_path: the empty path from the starting node to itself.
/// board: the initial Board. Must contain NodeStatus::default() for all Nodes.
/// get_neighbor: given a Node, returns an iterator of all its neighbor Edges
/// and Nodes.
/// report: receives a result of the shortest Path to a Node,
/// and returns whether to continue the algorithm.
pub fn dijkstra<
Node, Edge, Path, Board, GetNeighbor, NeighborIter, ReportResult
> (
start_node: Node,
start_path: Path,
mut board: Board,
mut get_neighbor: GetNeighbor,
mut report: ReportResult,
)
where
// ideally Board should be higher rank type over IndexMut::Output
// Board: for<'a, N: Default + Clone + Copy> IndexMut<&'a Node, Output = N>,
Board: for<'a> IndexMut<&'a Node, Output = NodeStatus>,
Path: Ord,
for<'a> &'a Path: Add<Edge, Output = Path>,
GetNeighbor: FnMut(Node) -> NeighborIter,
NeighborIter: IntoIterator<Item = (Edge, Node)>,
ReportResult: FnMut(&Node, &Path) -> bool,
{
let mut heap: Vec<(Path, Node)> = vec![];
board[&start_node] = NodeStatus::Visiting(0);
heap.push((start_path, start_node));
let swap_heap_element = |
heap: &mut Vec<(Path, Node)>,
board: &mut Board,
i: usize,
j: usize
| {
if i == j {
return;
}
board[&heap[i].1] = NodeStatus::Visiting(j);
board[&heap[j].1] = NodeStatus::Visiting(i);
heap.swap(i, j);
};
let pop_heap = |
heap: &mut Vec<(Path, Node)>,
board: &mut Board
| -> (Path, Node) {
swap_heap_element(heap, board, 0, heap.len() - 1);
let popped = heap.pop().unwrap();
if !heap.is_empty() {
let mut cur = 0;
loop {
let left = heap.get(cur * 2 + 1).map(|n|&n.0);
let right = heap.get(cur * 2 + 2).map(|n|&n.0);
if (left.is_none() || &heap[cur].0 <= left.unwrap())
&& (right.is_none() || &heap[cur].0 <= right.unwrap()) {
break;
}
let smaller =
if right.is_none() || left.unwrap() < right.unwrap() {
cur * 2 + 1
} else {
cur * 2 + 2
};
swap_heap_element(heap, board, cur, smaller);
cur = smaller;
}
}
board[&popped.1] = NodeStatus::Visited;
popped
};
let new_path = |
heap: &mut Vec<(Path, Node)>,
board: &mut Board,
path: Path,
node: Node
| {
let mut cur;
match board[&node] {
NodeStatus::Unknown => {
cur = heap.len();
board[&node] = NodeStatus::Visiting(cur);
heap.push((path, node));
}
NodeStatus::Visiting(i) => {
if heap[i].0 <= path {
return;
}
cur = i;
heap[i].0 = path;
}
NodeStatus::Visited => return
}
while cur != 0 {
let parent = (cur - 1) / 2;
if heap[parent].0 <= heap[cur].0 {
break;
}
swap_heap_element(heap, board, parent, cur);
cur = parent;
}
};
while !heap.is_empty() {
let (path, node) = pop_heap(&mut heap, &mut board);
if !report(&node, &path) {
return;
}
for (edge, new_node) in get_neighbor(node) {
new_path(&mut heap, &mut board, &path + edge, new_node);
}
}
}
/////////////////////////////////////////////////////////////////////////////
// Leetcode 407
enum Node {
Outside,
Grid(usize, usize)
}
struct Board {
width: usize,
grid: Vec<NodeStatus>,
outside: NodeStatus,
}
impl Board {
fn new(width: usize, height: usize) -> Board {
Board {
width,
grid: vec![NodeStatus::default(); width * height],
outside: NodeStatus::default()
}
}
}
#[derive(Copy, Clone, Default, PartialOrd, Ord, PartialEq, Eq)]
struct Path(i32);
impl<'a> Add<i32> for &'a Path {
type Output = Path;
fn add(self, other: i32) -> Path {
Path(std::cmp::max(self.0, other))
}
}
impl<'a> Index<&'a Node> for Board {
type Output = NodeStatus;
fn index(&self, index: &'a Node) -> &Self::Output {
match index {
Node::Outside => &self.outside,
Node::Grid(x, y) => &self.grid[*x + *y * self.width]
}
}
}
impl<'a> IndexMut<&'a Node> for Board {
fn index_mut(&mut self, index: &'a Node) -> &mut Self::Output {
match index {
Node::Outside => &mut self.outside,
Node::Grid(x, y) => &mut self.grid[*x + *y * self.width]
}
}
}
impl Solution {
pub fn trap_rain_water(height_map: Vec<Vec<i32>>) -> i32 {
let width = height_map.len();
let height = height_map[0].len();
let mut result = 0;
let get_neighbor = |node: Node| {
let mut list = vec![];
match node {
Node::Outside => {
for x in 0..width {
list.push((height_map[x][0], Node::Grid(x, 0)));
if height != 1 {
list.push((height_map[x][height - 1],
Node::Grid(x, height - 1)));
}
}
for y in 1..height - 1 {
list.push((height_map[0][y], Node::Grid(0, y)));
if width != 1 {
list.push((height_map[width - 1][y],
Node::Grid(width - 1, y)));
}
}
}
Node::Grid(x, y) => {
let mut push = |nx: usize, ny: usize| {
list.push((height_map[nx][ny], Node::Grid(nx, ny)));
};
if x > 0 {
push(x - 1, y);
}
if y > 0 {
push(x, y - 1);
}
if x < width - 1 {
push(x + 1, y);
}
if y < height - 1 {
push(x, y + 1);
}
}
}
list
};
let report = |node: &Node, path: &Path| -> bool {
if let Node::Grid(x, y) = node {
result += path.0 - height_map[*x][*y];
}
true
};
dijkstra(Node::Outside, Path::default(),
Board::new(width, height), get_neighbor, report
);
result
}
}
struct Solution;
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment