splay_tree/
splay-tree.rs

1use vec_arena::Arena;
2
3/// The null index, akin to null pointers.
4///
5/// Just like a null pointer indicates an address no object is ever stored at,
6/// the null index indicates an index no object is ever stored at.
7///
8/// Number `!0` is the largest possible value representable by `usize`.
9const NULL: usize = !0;
10
11struct Node<T> {
12    /// Parent node.
13    parent: usize,
14
15    /// Left and right child.
16    children: [usize; 2],
17
18    /// Actual value stored in node.
19    value: T,
20}
21
22impl<T> Node<T> {
23    fn new(value: T) -> Node<T> {
24        Node {
25            parent: NULL,
26            children: [NULL, NULL],
27            value: value,
28        }
29    }
30}
31
32struct Splay<T> {
33    /// This is where nodes are stored.
34    arena: Arena<Node<T>>,
35
36    /// The root node.
37    root: usize,
38}
39
40impl<T: Ord> Splay<T> {
41    /// Constructs a new, empty splay tree.
42    fn new() -> Splay<T> {
43        Splay {
44            arena: Arena::new(),
45            root: NULL,
46        }
47    }
48
49    /// Links nodes `p` and `c` as parent and child with the specified direction.
50    #[inline(always)]
51    fn link(&mut self, p: usize, c: usize, dir: usize) {
52        self.arena[p].children[dir] = c;
53        if c != NULL {
54            self.arena[c].parent = p;
55        }
56    }
57
58    /// Performs a rotation on node `c`, whose parent is node `p`.
59    #[inline(always)]
60    fn rotate(&mut self, p: usize, c: usize) {
61        // Variables:
62        // - `c` is the child node
63        // - `p` is it's parent
64        // - `g` is it's grandparent
65
66        // Find the grandparent.
67        let g = self.arena[p].parent;
68
69        // The direction of p-c relationship.
70        let dir = if self.arena[p].children[0] == c { 0 } else { 1 };
71
72        // This is the child of `c` that needs to be reassigned to `p`.
73        let t = self.arena[c].children[dir ^ 1];
74
75        self.link(p, t, dir);
76        self.link(c, p, dir ^ 1);
77
78        if g == NULL {
79            // There is no grandparent, so `c` becomes the root.
80            self.root = c;
81            self.arena[c].parent = NULL;
82        } else {
83            // Link `g` and `c` together.
84            let dir = if self.arena[g].children[0] == p { 0 } else { 1 };
85            self.link(g, c, dir);
86        }
87    }
88
89    /// Splays a node, rebalancing the tree in process.
90    fn splay(&mut self, c: usize) {
91        loop {
92            // Variables:
93            // - `c` is the current node
94            // - `p` is it's parent
95            // - `g` is it's grandparent
96
97            // Find the parent.
98            let p = self.arena[c].parent;
99            if p == NULL {
100                // There is no parent. That means `c` is the root.
101                break;
102            }
103
104            // Find the grandparent.
105            let g = self.arena[p].parent;
106            if g == NULL {
107                // There is no grandparent. Just one more rotation is left.
108                // Zig step.
109                self.rotate(p, c);
110                break;
111            }
112
113            if (self.arena[g].children[0] == p) == (self.arena[p].children[0] == c) {
114                // Zig-zig step.
115                self.rotate(g, p);
116                self.rotate(p, c);
117            } else {
118                // Zig-zag step.
119                self.rotate(p, c);
120                self.rotate(g, c);
121            }
122        }
123    }
124
125    /// Inserts a new node with specified `value`.
126    fn insert(&mut self, value: T) {
127        // Variables:
128        // - `n` is the new node
129        // - `p` will be it's parent
130        // - `c` is the present child of `p`
131
132        let n = self.arena.insert(Node::new(value));
133
134        if self.root == NULL {
135            self.root = n;
136        } else {
137            let mut p = self.root;
138            loop {
139                // Decide whether to go left or right.
140                let dir = if self.arena[n].value < self.arena[p].value {
141                    0
142                } else {
143                    1
144                };
145                let c = self.arena[p].children[dir];
146
147                if c == NULL {
148                    self.link(p, n, dir);
149                    self.splay(n);
150                    break;
151                }
152                p = c;
153            }
154        }
155    }
156
157    /// Pretty-prints the subtree rooted at `node`, indented by `indent` spaces.
158    fn print(&self, node: usize, indent: usize)
159    where
160        T: std::fmt::Display,
161    {
162        if node != NULL {
163            // Print the left subtree.
164            self.print(self.arena[node].children[0], indent + 3);
165
166            // Print the current node.
167            println!("{:width$}{}", "", self.arena[node].value, width = indent);
168
169            // Print the right subtree.
170            self.print(self.arena[node].children[1], indent + 3);
171        }
172    }
173}
174
175fn main() {
176    let mut splay = Splay::new();
177
178    // Insert a bunch of pseudorandom numbers.
179    let mut num = 1u32;
180    for _ in 0..30 {
181        num = num.wrapping_mul(17).wrapping_add(255);
182        splay.insert(num);
183    }
184
185    // Display the whole splay tree.
186    splay.print(splay.root, 0);
187}