class BSTIterator implements Iterator
public class Set implements Iterable
1: //Set.java
3: import java.lang.Iterable;
4: import java.util.Iterator;
5: import java.util.function.Function;
6: import java.util.function.Predicate;
8: class BSTIterator implements Iterator<Integer>
9: {
10: private BSTNode nextNode;
12: public BSTIterator(BSTNode startNode)
13: {
14: nextNode = startNode;
15: }
17: public boolean hasNext()
18: {
19: return nextNode != null;
20: }
22: public Integer next()
23: {
24: Integer returnMe = null;
25: if (nextNode != null)
26: {
27: returnMe = nextNode.data;
28: nextNode = nextNode.getSuccessor();
29: }
30: return returnMe;
31: }
32: }
34: public class Set implements Iterable<Integer>
35: {
36: private BSTNode root;
38: public Set()
39: {
40: root = null;
41: }
43: public boolean add(Integer newElement)
44: {
45: if (contains(newElement))
46: {
47: return false;
48: }
50: BSTNode newNode = new BSTNode(newElement, null);
52: // Special case for empty set
53: if (root == null)
54: {
55: root = newNode;
56: return true;
57: }
59: BSTNode node = root;
60: while (node != null)
61: {
62: if (newElement < node.data)
63: {
64: // Go left
65: if (node.left != null)
66: {
67: node = node.left;
68: }
69: else
70: {
71: node.left = newNode;
72: newNode.parent = node;
73: node = null;
74: }
75: }
76: else
77: {
78: // Go right
79: if (node.right != null)
80: {
81: node = node.right;
82: }
83: else
84: {
85: node.right = newNode;
86: newNode.parent = node;
87: node = null;
88: }
89: }
90: }
91: return true;
92: }
94: public boolean contains(Integer element)
95: {
96: return nodeSearch(element) != null;
97: }
99: public Set difference(Set otherSet)
100: {
101: Set result = new Set();
102: for (Integer element : this)
103: {
104: if (!otherSet.contains(element))
105: {
106: result.add(element);
107: }
108: }
109: return result;
110: }
112: public Set filter(Predicate<Integer> predicate)
113: {
114: Set result = new Set();
115: for (Integer element : this)
116: {
117: if (predicate.test(element))
118: {
119: result.add(element);
120: }
121: }
122: return result;
123: }
125: public Set intersection(Set otherSet)
126: {
127: Set result = new Set();
128: for (Integer element : this)
129: {
130: if (otherSet.contains(element))
131: {
132: result.add(element);
133: }
134: }
135: return result;
136: }
138: public Iterator<Integer> iterator()
139: {
140: // Special case for empty set
141: if (root == null)
142: {
143: return new BSTIterator(null);
144: }
146: // Start the iterator at the minimum node
147: BSTNode minNode = root;
148: while (minNode.left != null)
149: {
150: minNode = minNode.left;
151: }
152: return new BSTIterator(minNode);
153: }
155: public Set map(Function<Integer, Integer> mapFunction)
156: {
157: Set result = new Set();
158: for (Integer element : this)
159: {
160: Integer newElement = mapFunction.apply(element);
161: result.add(newElement);
162: }
163: return result;
164: }
166: private BSTNode nodeSearch(Integer element)
167: {
168: // Search the BST
169: BSTNode node = root;
170: while (node != null)
171: {
172: // Compare node's data against the search element
173: if (element.equals(node.data))
174: {
175: return node;
176: }
177: else if (element > node.data)
178: {
179: node = node.right;
180: }
181: else
182: {
183: node = node.left;
184: }
185: }
186: return node;
187: }
189: public void remove(Integer element)
190: {
191: removeNode(nodeSearch(element));
192: }
194: private void removeNode(BSTNode nodeToRemove)
195: {
196: if (nodeToRemove == null)
197: {
198: return;
199: }
201: // Case 1: Internal node with 2 children
202: if (nodeToRemove.left != null && nodeToRemove.right != null)
203: {
204: BSTNode successor = nodeToRemove.getSuccessor();
206: // Copy the data value from the successor
207: int dataCopy = successor.data;
209: // Remove successor
210: removeNode(successor);
212: // Replace nodeToRemove's data with successor data
213: nodeToRemove.data = dataCopy;
214: }
216: // Case 2: Root node (with 1 or 0 children)
217: else if (nodeToRemove == root)
218: {
219: if (nodeToRemove.left != null)
220: {
221: root = nodeToRemove.left;
222: }
223: else
224: {
225: root = nodeToRemove.right;
226: }
228: if (root != null)
229: {
230: root.parent = null;
231: }
232: }
234: // Case 3: Internal node with left child only
235: else if (nodeToRemove.left != null)
236: {
237: nodeToRemove.parent.replaceChild(nodeToRemove, nodeToRemove.left);
238: }
240: // Case 4: Internal node with right child only, or leaf node
241: else
242: {
243: nodeToRemove.parent.replaceChild(nodeToRemove, nodeToRemove.right);
244: }
245: }
247: public int size()
248: {
249: if (root == null)
250: {
251: return 0;
252: }
253: return root.count();
254: }
256: public Set union(Set otherSet)
257: {
258: Set result = new Set();
259: for (Integer element : this)
260: {
261: result.add(element);
262: }
263: for (Integer element : otherSet)
264: {
265: result.add(element);
266: }
267: return result;
268: }
269: }