Skip to content

Commit

Permalink
Add persistent tree in Rust (#197)
Browse files Browse the repository at this point in the history
  • Loading branch information
indy256 authored Jul 7, 2024
1 parent 724a847 commit 1127ff1
Show file tree
Hide file tree
Showing 2 changed files with 82 additions and 0 deletions.
4 changes: 4 additions & 0 deletions rust/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -31,3 +31,7 @@ path = "structures/disjoint_sets.rs"
[[bin]]
name = "fenwick_tree"
path = "structures/fenwick_tree.rs"

[[bin]]
name = "persistent_tree"
path = "structures/persistent_tree.rs"
78 changes: 78 additions & 0 deletions rust/structures/persistent_tree.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,78 @@
use std::rc::Rc;

pub struct Node {
left: Option<Rc<Node>>,
right: Option<Rc<Node>>,
sum: i32,
}

impl Node {
fn new(sum: i32) -> Self {
Self {
left: None,
right: None,
sum,
}
}

fn new_parent(left: &Rc<Node>, right: &Rc<Node>) -> Self {
Self {
left: Some(left.clone()),
right: Some(right.clone()),
sum: left.sum + right.sum,
}
}
}

pub struct PersistentTree;
impl PersistentTree {
pub fn build(left: i32, right: i32) -> Rc<Node> {
let node = if left == right {
Node::new(0)
} else {
let mid = (left + right) >> 1;
Node::new_parent(&Self::build(left, mid), &Self::build(mid + 1, right))
};
Rc::new(node)
}

pub fn sum(from: i32, to: i32, root: &Node, left: i32, right: i32) -> i32 {
if from > right || left > to {
0
} else if from <= left && right <= to {
root.sum
} else {
let mid = (left + right) >> 1;
Self::sum(from, to, root.left.as_ref().unwrap(), left, mid)
+ Self::sum(from, to, root.right.as_ref().unwrap(), mid + 1, right)
}
}

pub fn set(pos: i32, value: i32, root: &Node, left: i32, right: i32) -> Rc<Node> {
let node = if left == right {
Node::new(value)
} else {
let mid = (left + right) >> 1;
if pos <= mid {
Node::new_parent(&Self::set(pos, value, root.left.as_ref().unwrap(), left, mid), root.right.as_ref().unwrap())
} else {
Node::new_parent(root.left.as_ref().unwrap(), &Self::set(pos, value, root.right.as_ref().unwrap(), mid + 1, right))
}
};
Rc::new(node)
}
}

#[cfg(test)]
mod tests {
use crate::PersistentTree;

#[test]
fn basic_test() {
let n = 10;
let t1 = PersistentTree::build(0, n - 1);
let t2 = PersistentTree::set(0, 1, &t1, 0, n - 1);
assert_eq!(PersistentTree::sum(0, n - 1, &t1, 0, n - 1), 0);
assert_eq!(PersistentTree::sum(0, n - 1, &t2, 0, n - 1), 1);
}
}

0 comments on commit 1127ff1

Please sign in to comment.