From 1127ff16c4a8284433a44f0119f47490dd8ea968 Mon Sep 17 00:00:00 2001
From: Andrei Navumenka <indy256@gmail.com>
Date: Sun, 7 Jul 2024 13:43:37 -0400
Subject: [PATCH] Add persistent tree in Rust (#197)

---
 rust/Cargo.toml                    |  4 ++
 rust/structures/persistent_tree.rs | 78 ++++++++++++++++++++++++++++++
 2 files changed, 82 insertions(+)
 create mode 100644 rust/structures/persistent_tree.rs

diff --git a/rust/Cargo.toml b/rust/Cargo.toml
index ea28a970..ff242329 100644
--- a/rust/Cargo.toml
+++ b/rust/Cargo.toml
@@ -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"
diff --git a/rust/structures/persistent_tree.rs b/rust/structures/persistent_tree.rs
new file mode 100644
index 00000000..1f18f6a3
--- /dev/null
+++ b/rust/structures/persistent_tree.rs
@@ -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);
+    }
+}