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}