001package algs33;
002import stdlib.*;
003import algs13.Queue;
004/* ***********************************************************************
005 *  Compilation:  javac RedBlackLiteBST.java
006 *  Execution:    java RedBlackLiteBST < input.txt
007 *  Dependencies: StdIn.java StdOut.java
008 *  Data files:   http://algs4.cs.princeton.edu/33balanced/tinyST.txt
009 *
010 *  A symbol table implemented using a left-leaning red-black BST.
011 *  This is the 2-3 version.
012 *
013 *  This implementation implements only put, get, and contains.
014 *  See RedBlackBST.java for a full implementation including delete.
015 *
016 *
017 *  % more tinyST.txt
018 *  S E A R C H E X A M P L E
019 *
020 *  % java RedBlackLiteBST < tinyST.txt
021 *  A 8
022 *  C 4
023 *  E 12
024 *  H 5
025 *  L 11
026 *  M 9
027 *  P 10
028 *  R 3
029 *  S 0
030 *  X 7
031 *
032 *************************************************************************/
033
034public class XRedBlackLiteBST<K extends Comparable<? super K>, V> {
035
036        private static final boolean RED   = true;
037        private static final boolean BLACK = false;
038
039        private Node<K,V> root; // root of the BST
040        private int N;          // number of key-value pairs in BST
041
042        // BST helper node data type
043        private static class Node<K,V> {
044                public final K key;         // key
045                public V val;         // associated data
046                public Node<K,V> left, right;  // links to left and right subtrees
047                public boolean color;     // color of parent link
048
049                public Node(K key, V val, boolean color) {
050                        this.key = key;
051                        this.val = val;
052                        this.color = color;
053                }
054        }
055
056        /* ***********************************************************************
057         *  Standard BST search
058         *************************************************************************/
059
060        // return value associated with the given key, or null if no such key exists
061        public V get(K key) { return get(root, key); }
062        public V get(Node<K,V> x, K key) {
063                while (x != null) {
064                        int cmp = key.compareTo(x.key);
065                        if      (cmp < 0) x = x.left;
066                        else if (cmp > 0) x = x.right;
067                        else              return x.val;
068                }
069                return null;
070        }
071
072        // is there a key-value pair in the symbol table with the given key?
073        public boolean contains(K key) {
074                return (get(key) != null);
075        }
076
077
078        /* ***********************************************************************
079         *  Red-black insertion
080         *************************************************************************/
081
082        public void put(K key, V val) {
083                root = insert(root, key, val);
084                root.color = BLACK;
085                assert check();
086        }
087
088        private Node<K,V> insert(Node<K,V> h, K key, V val) {
089                if (h == null) {
090                        N++;
091                        return new Node<>(key, val, RED);
092                }
093
094                int cmp = key.compareTo(h.key);
095                if      (cmp < 0) h.left  = insert(h.left,  key, val);
096                else if (cmp > 0) h.right = insert(h.right, key, val);
097                else              h.val   = val;
098
099                // fix-up any right-leaning links
100                if (isRed(h.right) && !isRed(h.left))      h = rotateLeft(h);
101                if (isRed(h.left)  &&  isRed(h.left.left)) h = rotateRight(h);
102                if (isRed(h.left)  &&  isRed(h.right))     flipColors(h);
103
104                return h;
105        }
106
107        /* ***********************************************************************
108         *  red-black tree helper functions
109         *************************************************************************/
110
111        // is node x red (and non-null) ?
112        private boolean isRed(Node<K,V> x) {
113                if (x == null) return false;
114                return (x.color == RED);
115        }
116
117        // rotate right
118        private Node<K,V> rotateRight(Node<K,V> h) {
119                assert (h != null) && isRed(h.left);
120                Node<K,V> x = h.left;
121                h.left = x.right;
122                x.right = h;
123                x.color = h.color;
124                h.color = RED;
125                return x;
126        }
127
128        // rotate left
129        private Node<K,V> rotateLeft(Node<K,V> h) {
130                assert (h != null) && isRed(h.right);
131                Node<K,V> x = h.right;
132                h.right = x.left;
133                x.left = h;
134                x.color = h.color;
135                h.color = RED;
136                return x;
137        }
138
139        // precondition: two children are red, node is black
140        // postcondition: two children are black, node is red
141        private void flipColors(Node<K,V> h) {
142                assert !isRed(h) && isRed(h.left) && isRed(h.right);
143                h.color = RED;
144                h.left.color = BLACK;
145                h.right.color = BLACK;
146        }
147
148
149        /* ***********************************************************************
150         *  Utility functions
151         *************************************************************************/
152        // return number of key-value pairs in symbol table
153        public int size() { return N; }
154
155        // is the symbol table empty?
156        public boolean isEmpty() { return N == 0; }
157
158        // height of tree (empty tree height = 0)
159        public int height() { return height(root); }
160        private int height(Node<K,V> x) {
161                if (x == null) return 0;
162                return 1 + Math.max(height(x.left), height(x.right));
163        }
164
165        // return the smallest key; null if no such key
166        public K min() { return min(root); }
167        private K min(Node<K,V> x) {
168                K key = null;
169                while (x != null) {
170                        key = x.key;
171                        x = x.left;
172                }
173                return key;
174        }
175
176        // return the largest key; null if no such key
177        public K max() { return max(root); }
178        private K max(Node<K,V> x) {
179                K key = null;
180                while (x != null) {
181                        key = x.key;
182                        x = x.right;
183                }
184                return key;
185        }
186
187
188        /* *********************************************************************
189         *  Iterate using an inorder traversal.
190         *  Iterating through N elements takes O(N) time.
191         ***********************************************************************/
192        public Iterable<K> keys() {
193                Queue<K> queue = new Queue<>();
194                keys(root, queue);
195                return queue;
196        }
197
198        private void keys(Node<K,V> x, Queue<K> queue) {
199                if (x == null) return;
200                keys(x.left, queue);
201                queue.enqueue(x.key);
202                keys(x.right, queue);
203        }
204
205
206        /* ***********************************************************************
207         *  Check integrity of red-black BST data structure
208         *************************************************************************/
209        private boolean check() {
210                if (!isBST())            StdOut.println("Not in symmetric order");
211                if (!is23())             StdOut.println("Not a 2-3 tree");
212                if (!isBalanced())       StdOut.println("Not balanced");
213                return isBST() && is23() && isBalanced();
214        }
215
216        // does this binary tree satisfy symmetric order?
217        // Note: this test also ensures that data structure is a binary tree since order is strict
218        private boolean isBST() {
219                return isBST(root, null, null);
220        }
221
222        // is the tree rooted at x a BST with all keys strictly between min and max
223        // (if min or max is null, treat as empty constraint)
224        // Credit: Bob Dondero's elegant solution
225        private boolean isBST(Node<K,V> x, K min, K max) {
226                if (x == null) return true;
227                if (min != null && x.key.compareTo(min) <= 0) return false;
228                if (max != null && x.key.compareTo(max) >= 0) return false;
229                return isBST(x.left, min, x.key) && isBST(x.right, x.key, max);
230        }
231
232        // Does the tree have no red right links, and at most one (left)
233        // red links in a row on any path?
234        private boolean is23() { return is23(root); }
235        private boolean is23(Node<K,V> x) {
236                if (x == null) return true;
237                if (isRed(x.right)) return false;
238                if (x != root && isRed(x) && isRed(x.left))
239                        return false;
240                return is23(x.left) && is23(x.right);
241        }
242
243        // do all paths from root to leaf have same number of black edges?
244        private boolean isBalanced() {
245                int black = 0;     // number of black links on path from root to min
246                Node<K,V> x = root;
247                while (x != null) {
248                        if (!isRed(x)) black++;
249                        x = x.left;
250                }
251                return isBalanced(root, black);
252        }
253
254        // does every path from the root to a leaf have the given number of black links?
255        private boolean isBalanced(Node<K,V> x, int black) {
256                if (x == null) return black == 0;
257                if (!isRed(x)) black--;
258                return isBalanced(x.left, black) && isBalanced(x.right, black);
259        }
260
261
262        /* ***********************************************************************
263         *  Test client
264         *************************************************************************/
265        public static void main(String[] args) {
266
267                String test = "S E A R C H E X A M P L E";
268                String[] keys = test.split(" ");
269                XRedBlackLiteBST<String, Integer> st = new XRedBlackLiteBST<>();
270                for (int i = 0; i < keys.length; i++)
271                        st.put(keys[i], i);
272
273                StdOut.println("size = " + st.size());
274                StdOut.println("min  = " + st.min());
275                StdOut.println("max  = " + st.max());
276                StdOut.println();
277
278
279                // print keys in order using allKeys()
280                StdOut.println("Testing keys()");
281                StdOut.println("--------------------------------");
282                for (String s : st.keys())
283                        StdOut.println(s + " " + st.get(s));
284                StdOut.println();
285
286                // insert N elements in order if one command-line argument supplied
287                if (args.length == 0) return;
288                int N = Integer.parseInt(args[0]);
289                XRedBlackLiteBST<Integer, Integer> st2 = new XRedBlackLiteBST<>();
290                for (int i = 0; i < N; i++) {
291                        st2.put(i, i);
292                        int h = st2.height();
293                        StdOut.println("i = " + i + ", height = " + h + ", size = " + st2.size());
294                }
295
296
297                StdOut.println("size = " + st2.size());
298        }
299}