diff --git a/solution/2400-2499/2454.Next Greater Element IV/README.md b/solution/2400-2499/2454.Next Greater Element IV/README.md index 05da24c904279..caf0a7a09dfd8 100644 --- a/solution/2400-2499/2454.Next Greater Element IV/README.md +++ b/solution/2400-2499/2454.Next Greater Element IV/README.md @@ -62,15 +62,13 @@ -**方法一:单调栈 + 优先队列(小根堆)** +**方法一:排序 + 有序集合** -求下一个更大的元素,可以使用单调栈来实现。我们维护一个栈,然后从左到右遍历数组,如果栈顶元素小于当前元素,则当前元素就是栈顶元素的下一个更大的元素。 +我们可以将数组中的元素转成二元组 $(x, i)$,其中 $x$ 为元素的值,$i$ 为元素的下标。然后按照元素的值从大到小排序。 -这道题的变形是求下一个更大的元素的下一个更大的元素,即第二大的元素。我们观察单调栈求下一个更大元素的过程,每次出栈时,栈顶元素找到了下一个更大的元素,但我们是要为栈顶元素找到第二个更大的元素。次数,我们可以将栈顶元素出栈,放到一个优先队列(小根堆)中。每次遍历数组元素时,先判断当前元素是否大于优先队列的堆顶元素,如果大于,说明堆顶元素找打了第二个更大的元素,更新答案数组,然后弹出堆顶元素,继续判断当前元素是否大于优先队列的堆顶元素,直到堆为空或者当前元素不大于堆顶元素。 +接下来,我们遍历排序后的数组,维护一个有序集合,其中存储的是元素的下标。当我们遍历到元素 $(x, i)$ 时,所有大于 $x$ 的元素的下标都已经在有序集合中了。我们只需要在有序集合中,找到 $i$ 的下下一个元素的下标 $j$,那么 $j$ 对应的元素就是 $x$ 的第二大元素。然后,我们将 $i$ 加入到有序集合中。继续遍历下一个元素。 -接着,执行单调栈的相关操作,弹出栈顶元素后,放入到优先队列中。 - -时间复杂度 $O(n\log n)$,空间复杂度 $O(n)$。其中 $n$ 为数组的长度。 +时间复杂度 $O(n \times \log n)$,空间复杂度 $O(n)$。其中 $n$ 为数组的长度。 @@ -79,18 +77,21 @@ ```python +from sortedcontainers import SortedList + + class Solution: def secondGreaterElement(self, nums: List[int]) -> List[int]: - stk = [] - q = [] - ans = [-1] * len(nums) - for i, v in enumerate(nums): - while q and q[0][0] < v: - ans[q[0][1]] = v - heappop(q) - while stk and nums[stk[-1]] < v: - heappush(q, (nums[stk[-1]], stk.pop())) - stk.append(i) + arr = [(x, i) for i, x in enumerate(nums)] + arr.sort(key=lambda x: -x[0]) + sl = SortedList() + n = len(nums) + ans = [-1] * n + for _, i in arr: + j = sl.bisect_right(i) + if j + 1 < len(sl): + ans[i] = nums[sl[j + 1]] + sl.add(i) return ans ``` @@ -101,21 +102,22 @@ class Solution: ```java class Solution { public int[] secondGreaterElement(int[] nums) { - Deque stk = new ArrayDeque<>(); - PriorityQueue q = new PriorityQueue<>((a, b) -> a[0] - b[0]); int n = nums.length; int[] ans = new int[n]; Arrays.fill(ans, -1); + int[][] arr = new int[n][0]; for (int i = 0; i < n; ++i) { - int v = nums[i]; - while (!q.isEmpty() && q.peek()[0] < v) { - ans[q.peek()[1]] = v; - q.poll(); - } - while (!stk.isEmpty() && nums[stk.peek()] < v) { - q.offer(new int[] {nums[stk.peek()], stk.pop()}); + arr[i] = new int[] {nums[i], i}; + } + Arrays.sort(arr, (a, b) -> b[0] - a[0]); + TreeSet ts = new TreeSet<>(); + for (int[] pair : arr) { + int i = pair[1]; + Integer j = ts.higher(i); + if (j != null && ts.higher(j) != null) { + ans[i] = nums[ts.higher(j)]; } - stk.push(i); + ts.add(i); } return ans; } @@ -125,75 +127,693 @@ class Solution { ### **C++** ```cpp -using pii = pair; - class Solution { public: vector secondGreaterElement(vector& nums) { - stack stk; - priority_queue, greater> q; int n = nums.size(); vector ans(n, -1); + vector> arr(n); for (int i = 0; i < n; ++i) { - int v = nums[i]; - while (!q.empty() && q.top().first < v) { - ans[q.top().second] = v; - q.pop(); - } - while (!stk.empty() && nums[stk.top()] < v) { - q.push({nums[stk.top()], stk.top()}); - stk.pop(); + arr[i] = {-nums[i], i}; + } + sort(arr.begin(), arr.end()); + set ts; + for (auto& [_, i] : arr) { + auto it = ts.upper_bound(i); + if (it != ts.end() && ts.upper_bound(*it) != ts.end()) { + ans[i] = nums[*ts.upper_bound(*it)]; } - stk.push(i); + ts.insert(i); } return ans; } }; ``` -### **Go** - -```go -func secondGreaterElement(nums []int) []int { - stk := []int{} - q := hp{} - n := len(nums) - ans := make([]int, n) - for i := range ans { - ans[i] = -1 - } - for i, v := range nums { - for len(q) > 0 && q[0].v < v { - ans[q[0].i] = v - heap.Pop(&q) - } - for len(stk) > 0 && nums[stk[len(stk)-1]] < v { - heap.Push(&q, pair{nums[stk[len(stk)-1]], stk[len(stk)-1]}) - stk = stk[:len(stk)-1] - } - stk = append(stk, i) - } - return ans +### **TypeScript** + +```ts +function secondGreaterElement(nums: number[]): number[] { + const n = nums.length; + const arr: number[][] = []; + for (let i = 0; i < n; ++i) { + arr.push([nums[i], i]); + } + arr.sort((a, b) => b[0] - a[0]); + const ans = Array(n).fill(-1); + const ts = new TreeSet(); + for (const [_, i] of arr) { + let j = ts.higher(i); + if (j !== undefined) { + j = ts.higher(j); + if (j !== undefined) { + ans[i] = nums[j]; + } + } + ts.add(i); + } + return ans; } -type pair struct{ v, i int } +type Compare = (lhs: T, rhs: T) => number; + +class RBTreeNode { + data: T; + count: number; + left: RBTreeNode | null; + right: RBTreeNode | null; + parent: RBTreeNode | null; + color: number; + constructor(data: T) { + this.data = data; + this.left = this.right = this.parent = null; + this.color = 0; + this.count = 1; + } + + sibling(): RBTreeNode | null { + if (!this.parent) return null; // sibling null if no parent + return this.isOnLeft() ? this.parent.right : this.parent.left; + } -type hp []pair + isOnLeft(): boolean { + return this === this.parent!.left; + } -func (h hp) Len() int { return len(h) } -func (h hp) Less(i, j int) bool { - a, b := h[i], h[j] - return a.v < b.v + hasRedChild(): boolean { + return ( + Boolean(this.left && this.left.color === 0) || + Boolean(this.right && this.right.color === 0) + ); + } } -func (h hp) Swap(i, j int) { h[i], h[j] = h[j], h[i] } -func (h *hp) Push(v any) { *h = append(*h, v.(pair)) } -func (h *hp) Pop() any { a := *h; v := a[len(a)-1]; *h = a[:len(a)-1]; return v } -``` -### **TypeScript** +class RBTree { + root: RBTreeNode | null; + lt: (l: T, r: T) => boolean; + constructor(compare: Compare = (l: T, r: T) => (l < r ? -1 : l > r ? 1 : 0)) { + this.root = null; + this.lt = (l: T, r: T) => compare(l, r) < 0; + } -```ts + rotateLeft(pt: RBTreeNode): void { + const right = pt.right!; + pt.right = right.left; + + if (pt.right) pt.right.parent = pt; + right.parent = pt.parent; + + if (!pt.parent) this.root = right; + else if (pt === pt.parent.left) pt.parent.left = right; + else pt.parent.right = right; + + right.left = pt; + pt.parent = right; + } + + rotateRight(pt: RBTreeNode): void { + const left = pt.left!; + pt.left = left.right; + + if (pt.left) pt.left.parent = pt; + left.parent = pt.parent; + + if (!pt.parent) this.root = left; + else if (pt === pt.parent.left) pt.parent.left = left; + else pt.parent.right = left; + + left.right = pt; + pt.parent = left; + } + + swapColor(p1: RBTreeNode, p2: RBTreeNode): void { + const tmp = p1.color; + p1.color = p2.color; + p2.color = tmp; + } + + swapData(p1: RBTreeNode, p2: RBTreeNode): void { + const tmp = p1.data; + p1.data = p2.data; + p2.data = tmp; + } + + fixAfterInsert(pt: RBTreeNode): void { + let parent = null; + let grandParent = null; + + while (pt !== this.root && pt.color !== 1 && pt.parent?.color === 0) { + parent = pt.parent; + grandParent = pt.parent.parent; + + /* Case : A + Parent of pt is left child of Grand-parent of pt */ + if (parent === grandParent?.left) { + const uncle = grandParent.right; + + /* Case : 1 + The uncle of pt is also red + Only Recoloring required */ + if (uncle && uncle.color === 0) { + grandParent.color = 0; + parent.color = 1; + uncle.color = 1; + pt = grandParent; + } else { + /* Case : 2 + pt is right child of its parent + Left-rotation required */ + if (pt === parent.right) { + this.rotateLeft(parent); + pt = parent; + parent = pt.parent; + } + + /* Case : 3 + pt is left child of its parent + Right-rotation required */ + this.rotateRight(grandParent); + this.swapColor(parent!, grandParent); + pt = parent!; + } + } else { + /* Case : B + Parent of pt is right child of Grand-parent of pt */ + const uncle = grandParent!.left; + + /* Case : 1 + The uncle of pt is also red + Only Recoloring required */ + if (uncle != null && uncle.color === 0) { + grandParent!.color = 0; + parent.color = 1; + uncle.color = 1; + pt = grandParent!; + } else { + /* Case : 2 + pt is left child of its parent + Right-rotation required */ + if (pt === parent.left) { + this.rotateRight(parent); + pt = parent; + parent = pt.parent; + } + + /* Case : 3 + pt is right child of its parent + Left-rotation required */ + this.rotateLeft(grandParent!); + this.swapColor(parent!, grandParent!); + pt = parent!; + } + } + } + this.root!.color = 1; + } + delete(val: T): boolean { + const node = this.find(val); + if (!node) return false; + node.count--; + if (!node.count) this.deleteNode(node); + return true; + } + + deleteAll(val: T): boolean { + const node = this.find(val); + if (!node) return false; + this.deleteNode(node); + return true; + } + + deleteNode(v: RBTreeNode): void { + const u = BSTreplace(v); + + // True when u and v are both black + const uvBlack = (u === null || u.color === 1) && v.color === 1; + const parent = v.parent!; + + if (!u) { + // u is null therefore v is leaf + if (v === this.root) this.root = null; + // v is root, making root null + else { + if (uvBlack) { + // u and v both black + // v is leaf, fix double black at v + this.fixDoubleBlack(v); + } else { + // u or v is red + if (v.sibling()) { + // sibling is not null, make it red" + v.sibling()!.color = 0; + } + } + // delete v from the tree + if (v.isOnLeft()) parent.left = null; + else parent.right = null; + } + return; + } + + if (!v.left || !v.right) { + // v has 1 child + if (v === this.root) { + // v is root, assign the value of u to v, and delete u + v.data = u.data; + v.left = v.right = null; + } else { + // Detach v from tree and move u up + if (v.isOnLeft()) parent.left = u; + else parent.right = u; + u.parent = parent; + if (uvBlack) this.fixDoubleBlack(u); + // u and v both black, fix double black at u + else u.color = 1; // u or v red, color u black + } + return; + } + + // v has 2 children, swap data with successor and recurse + this.swapData(u, v); + this.deleteNode(u); + + // find node that replaces a deleted node in BST + function BSTreplace(x: RBTreeNode): RBTreeNode | null { + // when node have 2 children + if (x.left && x.right) return successor(x.right); + // when leaf + if (!x.left && !x.right) return null; + // when single child + return x.left ?? x.right; + } + // find node that do not have a left child + // in the subtree of the given node + function successor(x: RBTreeNode): RBTreeNode { + let temp = x; + while (temp.left) temp = temp.left; + return temp; + } + } + + fixDoubleBlack(x: RBTreeNode): void { + if (x === this.root) return; // Reached root + + const sibling = x.sibling(); + const parent = x.parent!; + if (!sibling) { + // No sibiling, double black pushed up + this.fixDoubleBlack(parent); + } else { + if (sibling.color === 0) { + // Sibling red + parent.color = 0; + sibling.color = 1; + if (sibling.isOnLeft()) this.rotateRight(parent); + // left case + else this.rotateLeft(parent); // right case + this.fixDoubleBlack(x); + } else { + // Sibling black + if (sibling.hasRedChild()) { + // at least 1 red children + if (sibling.left && sibling.left.color === 0) { + if (sibling.isOnLeft()) { + // left left + sibling.left.color = sibling.color; + sibling.color = parent.color; + this.rotateRight(parent); + } else { + // right left + sibling.left.color = parent.color; + this.rotateRight(sibling); + this.rotateLeft(parent); + } + } else { + if (sibling.isOnLeft()) { + // left right + sibling.right!.color = parent.color; + this.rotateLeft(sibling); + this.rotateRight(parent); + } else { + // right right + sibling.right!.color = sibling.color; + sibling.color = parent.color; + this.rotateLeft(parent); + } + } + parent.color = 1; + } else { + // 2 black children + sibling.color = 0; + if (parent.color === 1) this.fixDoubleBlack(parent); + else parent.color = 1; + } + } + } + } + + insert(data: T): boolean { + // search for a position to insert + let parent = this.root; + while (parent) { + if (this.lt(data, parent.data)) { + if (!parent.left) break; + else parent = parent.left; + } else if (this.lt(parent.data, data)) { + if (!parent.right) break; + else parent = parent.right; + } else break; + } + + // insert node into parent + const node = new RBTreeNode(data); + if (!parent) this.root = node; + else if (this.lt(node.data, parent.data)) parent.left = node; + else if (this.lt(parent.data, node.data)) parent.right = node; + else { + parent.count++; + return false; + } + node.parent = parent; + this.fixAfterInsert(node); + return true; + } + + find(data: T): RBTreeNode | null { + let p = this.root; + while (p) { + if (this.lt(data, p.data)) { + p = p.left; + } else if (this.lt(p.data, data)) { + p = p.right; + } else break; + } + return p ?? null; + } + + *inOrder(root: RBTreeNode = this.root!): Generator { + if (!root) return; + for (const v of this.inOrder(root.left!)) yield v; + yield root.data; + for (const v of this.inOrder(root.right!)) yield v; + } + + *reverseInOrder(root: RBTreeNode = this.root!): Generator { + if (!root) return; + for (const v of this.reverseInOrder(root.right!)) yield v; + yield root.data; + for (const v of this.reverseInOrder(root.left!)) yield v; + } +} + +class TreeSet { + _size: number; + tree: RBTree; + compare: Compare; + constructor( + collection: T[] | Compare = [], + compare: Compare = (l: T, r: T) => (l < r ? -1 : l > r ? 1 : 0), + ) { + if (typeof collection === 'function') { + compare = collection; + collection = []; + } + this._size = 0; + this.compare = compare; + this.tree = new RBTree(compare); + for (const val of collection) this.add(val); + } + + size(): number { + return this._size; + } + + has(val: T): boolean { + return !!this.tree.find(val); + } + + add(val: T): boolean { + const successful = this.tree.insert(val); + this._size += successful ? 1 : 0; + return successful; + } + + delete(val: T): boolean { + const deleted = this.tree.deleteAll(val); + this._size -= deleted ? 1 : 0; + return deleted; + } + + ceil(val: T): T | undefined { + let p = this.tree.root; + let higher = null; + while (p) { + if (this.compare(p.data, val) >= 0) { + higher = p; + p = p.left; + } else { + p = p.right; + } + } + return higher?.data; + } + + floor(val: T): T | undefined { + let p = this.tree.root; + let lower = null; + while (p) { + if (this.compare(val, p.data) >= 0) { + lower = p; + p = p.right; + } else { + p = p.left; + } + } + return lower?.data; + } + + higher(val: T): T | undefined { + let p = this.tree.root; + let higher = null; + while (p) { + if (this.compare(val, p.data) < 0) { + higher = p; + p = p.left; + } else { + p = p.right; + } + } + return higher?.data; + } + + lower(val: T): T | undefined { + let p = this.tree.root; + let lower = null; + while (p) { + if (this.compare(p.data, val) < 0) { + lower = p; + p = p.right; + } else { + p = p.left; + } + } + return lower?.data; + } + + first(): T | undefined { + return this.tree.inOrder().next().value; + } + + last(): T | undefined { + return this.tree.reverseInOrder().next().value; + } + + shift(): T | undefined { + const first = this.first(); + if (first === undefined) return undefined; + this.delete(first); + return first; + } + + pop(): T | undefined { + const last = this.last(); + if (last === undefined) return undefined; + this.delete(last); + return last; + } + + *[Symbol.iterator](): Generator { + for (const val of this.values()) yield val; + } + + *keys(): Generator { + for (const val of this.values()) yield val; + } + + *values(): Generator { + for (const val of this.tree.inOrder()) yield val; + return undefined; + } + + /** + * Return a generator for reverse order traversing the set + */ + *rvalues(): Generator { + for (const val of this.tree.reverseInOrder()) yield val; + return undefined; + } +} + +class TreeMultiSet { + _size: number; + tree: RBTree; + compare: Compare; + constructor( + collection: T[] | Compare = [], + compare: Compare = (l: T, r: T) => (l < r ? -1 : l > r ? 1 : 0), + ) { + if (typeof collection === 'function') { + compare = collection; + collection = []; + } + this._size = 0; + this.compare = compare; + this.tree = new RBTree(compare); + for (const val of collection) this.add(val); + } + + size(): number { + return this._size; + } + + has(val: T): boolean { + return !!this.tree.find(val); + } + + add(val: T): boolean { + const successful = this.tree.insert(val); + this._size++; + return successful; + } + + delete(val: T): boolean { + const successful = this.tree.delete(val); + if (!successful) return false; + this._size--; + return true; + } + + count(val: T): number { + const node = this.tree.find(val); + return node ? node.count : 0; + } + + ceil(val: T): T | undefined { + let p = this.tree.root; + let higher = null; + while (p) { + if (this.compare(p.data, val) >= 0) { + higher = p; + p = p.left; + } else { + p = p.right; + } + } + return higher?.data; + } + + floor(val: T): T | undefined { + let p = this.tree.root; + let lower = null; + while (p) { + if (this.compare(val, p.data) >= 0) { + lower = p; + p = p.right; + } else { + p = p.left; + } + } + return lower?.data; + } + + higher(val: T): T | undefined { + let p = this.tree.root; + let higher = null; + while (p) { + if (this.compare(val, p.data) < 0) { + higher = p; + p = p.left; + } else { + p = p.right; + } + } + return higher?.data; + } + + lower(val: T): T | undefined { + let p = this.tree.root; + let lower = null; + while (p) { + if (this.compare(p.data, val) < 0) { + lower = p; + p = p.right; + } else { + p = p.left; + } + } + return lower?.data; + } + + first(): T | undefined { + return this.tree.inOrder().next().value; + } + + last(): T | undefined { + return this.tree.reverseInOrder().next().value; + } + + shift(): T | undefined { + const first = this.first(); + if (first === undefined) return undefined; + this.delete(first); + return first; + } + + pop(): T | undefined { + const last = this.last(); + if (last === undefined) return undefined; + this.delete(last); + return last; + } + + *[Symbol.iterator](): Generator { + yield* this.values(); + } + + *keys(): Generator { + for (const val of this.values()) yield val; + } + + *values(): Generator { + for (const val of this.tree.inOrder()) { + let count = this.count(val); + while (count--) yield val; + } + return undefined; + } + + /** + * Return a generator for reverse order traversing the multi-set + */ + *rvalues(): Generator { + for (const val of this.tree.reverseInOrder()) { + let count = this.count(val); + while (count--) yield val; + } + return undefined; + } +} ``` ### **...** diff --git a/solution/2400-2499/2454.Next Greater Element IV/README_EN.md b/solution/2400-2499/2454.Next Greater Element IV/README_EN.md index a7d5aae768b33..5fcd54968e6be 100644 --- a/solution/2400-2499/2454.Next Greater Element IV/README_EN.md +++ b/solution/2400-2499/2454.Next Greater Element IV/README_EN.md @@ -56,23 +56,34 @@ We return [-1,-1] since neither integer has any integer greater than it. ## Solutions +**Solution 1: Sorting + Ordered Set** + +We can convert the elements in the array into pairs $(x, i)$, where $x$ is the value of the element and $i$ is the index of the element. Then sort by the value of the elements in descending order. + +Next, we traverse the sorted array, maintaining an ordered set that stores the indices of the elements. When we traverse to the element $(x, i)$, the indices of all elements greater than $x$ are already in the ordered set. We only need to find the index $j$ of the next element after $i$ in the ordered set, then the element corresponding to $j$ is the second largest element of $x$. Then, we add $i$ to the ordered set. Continue to traverse the next element. + +The time complexity is $O(n \times \log n)$, and the space complexity is $O(n)$. Here, $n$ is the length of the array. + ### **Python3** ```python +from sortedcontainers import SortedList + + class Solution: def secondGreaterElement(self, nums: List[int]) -> List[int]: - stk = [] - q = [] - ans = [-1] * len(nums) - for i, v in enumerate(nums): - while q and q[0][0] < v: - ans[q[0][1]] = v - heappop(q) - while stk and nums[stk[-1]] < v: - heappush(q, (nums[stk[-1]], stk.pop())) - stk.append(i) + arr = [(x, i) for i, x in enumerate(nums)] + arr.sort(key=lambda x: -x[0]) + sl = SortedList() + n = len(nums) + ans = [-1] * n + for _, i in arr: + j = sl.bisect_right(i) + if j + 1 < len(sl): + ans[i] = nums[sl[j + 1]] + sl.add(i) return ans ``` @@ -81,21 +92,22 @@ class Solution: ```java class Solution { public int[] secondGreaterElement(int[] nums) { - Deque stk = new ArrayDeque<>(); - PriorityQueue q = new PriorityQueue<>((a, b) -> a[0] - b[0]); int n = nums.length; int[] ans = new int[n]; Arrays.fill(ans, -1); + int[][] arr = new int[n][0]; for (int i = 0; i < n; ++i) { - int v = nums[i]; - while (!q.isEmpty() && q.peek()[0] < v) { - ans[q.peek()[1]] = v; - q.poll(); - } - while (!stk.isEmpty() && nums[stk.peek()] < v) { - q.offer(new int[] {nums[stk.peek()], stk.pop()}); + arr[i] = new int[] {nums[i], i}; + } + Arrays.sort(arr, (a, b) -> b[0] - a[0]); + TreeSet ts = new TreeSet<>(); + for (int[] pair : arr) { + int i = pair[1]; + Integer j = ts.higher(i); + if (j != null && ts.higher(j) != null) { + ans[i] = nums[ts.higher(j)]; } - stk.push(i); + ts.add(i); } return ans; } @@ -105,75 +117,693 @@ class Solution { ### **C++** ```cpp -using pii = pair; - class Solution { public: vector secondGreaterElement(vector& nums) { - stack stk; - priority_queue, greater> q; int n = nums.size(); vector ans(n, -1); + vector> arr(n); for (int i = 0; i < n; ++i) { - int v = nums[i]; - while (!q.empty() && q.top().first < v) { - ans[q.top().second] = v; - q.pop(); - } - while (!stk.empty() && nums[stk.top()] < v) { - q.push({nums[stk.top()], stk.top()}); - stk.pop(); + arr[i] = {-nums[i], i}; + } + sort(arr.begin(), arr.end()); + set ts; + for (auto& [_, i] : arr) { + auto it = ts.upper_bound(i); + if (it != ts.end() && ts.upper_bound(*it) != ts.end()) { + ans[i] = nums[*ts.upper_bound(*it)]; } - stk.push(i); + ts.insert(i); } return ans; } }; ``` -### **Go** - -```go -func secondGreaterElement(nums []int) []int { - stk := []int{} - q := hp{} - n := len(nums) - ans := make([]int, n) - for i := range ans { - ans[i] = -1 - } - for i, v := range nums { - for len(q) > 0 && q[0].v < v { - ans[q[0].i] = v - heap.Pop(&q) - } - for len(stk) > 0 && nums[stk[len(stk)-1]] < v { - heap.Push(&q, pair{nums[stk[len(stk)-1]], stk[len(stk)-1]}) - stk = stk[:len(stk)-1] - } - stk = append(stk, i) - } - return ans +### **TypeScript** + +```ts +function secondGreaterElement(nums: number[]): number[] { + const n = nums.length; + const arr: number[][] = []; + for (let i = 0; i < n; ++i) { + arr.push([nums[i], i]); + } + arr.sort((a, b) => b[0] - a[0]); + const ans = Array(n).fill(-1); + const ts = new TreeSet(); + for (const [_, i] of arr) { + let j = ts.higher(i); + if (j !== undefined) { + j = ts.higher(j); + if (j !== undefined) { + ans[i] = nums[j]; + } + } + ts.add(i); + } + return ans; } -type pair struct{ v, i int } +type Compare = (lhs: T, rhs: T) => number; + +class RBTreeNode { + data: T; + count: number; + left: RBTreeNode | null; + right: RBTreeNode | null; + parent: RBTreeNode | null; + color: number; + constructor(data: T) { + this.data = data; + this.left = this.right = this.parent = null; + this.color = 0; + this.count = 1; + } + + sibling(): RBTreeNode | null { + if (!this.parent) return null; // sibling null if no parent + return this.isOnLeft() ? this.parent.right : this.parent.left; + } -type hp []pair + isOnLeft(): boolean { + return this === this.parent!.left; + } -func (h hp) Len() int { return len(h) } -func (h hp) Less(i, j int) bool { - a, b := h[i], h[j] - return a.v < b.v + hasRedChild(): boolean { + return ( + Boolean(this.left && this.left.color === 0) || + Boolean(this.right && this.right.color === 0) + ); + } } -func (h hp) Swap(i, j int) { h[i], h[j] = h[j], h[i] } -func (h *hp) Push(v any) { *h = append(*h, v.(pair)) } -func (h *hp) Pop() any { a := *h; v := a[len(a)-1]; *h = a[:len(a)-1]; return v } -``` -### **TypeScript** +class RBTree { + root: RBTreeNode | null; + lt: (l: T, r: T) => boolean; + constructor(compare: Compare = (l: T, r: T) => (l < r ? -1 : l > r ? 1 : 0)) { + this.root = null; + this.lt = (l: T, r: T) => compare(l, r) < 0; + } -```ts + rotateLeft(pt: RBTreeNode): void { + const right = pt.right!; + pt.right = right.left; + + if (pt.right) pt.right.parent = pt; + right.parent = pt.parent; + + if (!pt.parent) this.root = right; + else if (pt === pt.parent.left) pt.parent.left = right; + else pt.parent.right = right; + + right.left = pt; + pt.parent = right; + } + + rotateRight(pt: RBTreeNode): void { + const left = pt.left!; + pt.left = left.right; + + if (pt.left) pt.left.parent = pt; + left.parent = pt.parent; + + if (!pt.parent) this.root = left; + else if (pt === pt.parent.left) pt.parent.left = left; + else pt.parent.right = left; + + left.right = pt; + pt.parent = left; + } + + swapColor(p1: RBTreeNode, p2: RBTreeNode): void { + const tmp = p1.color; + p1.color = p2.color; + p2.color = tmp; + } + + swapData(p1: RBTreeNode, p2: RBTreeNode): void { + const tmp = p1.data; + p1.data = p2.data; + p2.data = tmp; + } + + fixAfterInsert(pt: RBTreeNode): void { + let parent = null; + let grandParent = null; + + while (pt !== this.root && pt.color !== 1 && pt.parent?.color === 0) { + parent = pt.parent; + grandParent = pt.parent.parent; + + /* Case : A + Parent of pt is left child of Grand-parent of pt */ + if (parent === grandParent?.left) { + const uncle = grandParent.right; + + /* Case : 1 + The uncle of pt is also red + Only Recoloring required */ + if (uncle && uncle.color === 0) { + grandParent.color = 0; + parent.color = 1; + uncle.color = 1; + pt = grandParent; + } else { + /* Case : 2 + pt is right child of its parent + Left-rotation required */ + if (pt === parent.right) { + this.rotateLeft(parent); + pt = parent; + parent = pt.parent; + } + + /* Case : 3 + pt is left child of its parent + Right-rotation required */ + this.rotateRight(grandParent); + this.swapColor(parent!, grandParent); + pt = parent!; + } + } else { + /* Case : B + Parent of pt is right child of Grand-parent of pt */ + const uncle = grandParent!.left; + + /* Case : 1 + The uncle of pt is also red + Only Recoloring required */ + if (uncle != null && uncle.color === 0) { + grandParent!.color = 0; + parent.color = 1; + uncle.color = 1; + pt = grandParent!; + } else { + /* Case : 2 + pt is left child of its parent + Right-rotation required */ + if (pt === parent.left) { + this.rotateRight(parent); + pt = parent; + parent = pt.parent; + } + + /* Case : 3 + pt is right child of its parent + Left-rotation required */ + this.rotateLeft(grandParent!); + this.swapColor(parent!, grandParent!); + pt = parent!; + } + } + } + this.root!.color = 1; + } + + delete(val: T): boolean { + const node = this.find(val); + if (!node) return false; + node.count--; + if (!node.count) this.deleteNode(node); + return true; + } + + deleteAll(val: T): boolean { + const node = this.find(val); + if (!node) return false; + this.deleteNode(node); + return true; + } + + deleteNode(v: RBTreeNode): void { + const u = BSTreplace(v); + + // True when u and v are both black + const uvBlack = (u === null || u.color === 1) && v.color === 1; + const parent = v.parent!; + + if (!u) { + // u is null therefore v is leaf + if (v === this.root) this.root = null; + // v is root, making root null + else { + if (uvBlack) { + // u and v both black + // v is leaf, fix double black at v + this.fixDoubleBlack(v); + } else { + // u or v is red + if (v.sibling()) { + // sibling is not null, make it red" + v.sibling()!.color = 0; + } + } + // delete v from the tree + if (v.isOnLeft()) parent.left = null; + else parent.right = null; + } + return; + } + + if (!v.left || !v.right) { + // v has 1 child + if (v === this.root) { + // v is root, assign the value of u to v, and delete u + v.data = u.data; + v.left = v.right = null; + } else { + // Detach v from tree and move u up + if (v.isOnLeft()) parent.left = u; + else parent.right = u; + u.parent = parent; + if (uvBlack) this.fixDoubleBlack(u); + // u and v both black, fix double black at u + else u.color = 1; // u or v red, color u black + } + return; + } + + // v has 2 children, swap data with successor and recurse + this.swapData(u, v); + this.deleteNode(u); + + // find node that replaces a deleted node in BST + function BSTreplace(x: RBTreeNode): RBTreeNode | null { + // when node have 2 children + if (x.left && x.right) return successor(x.right); + // when leaf + if (!x.left && !x.right) return null; + // when single child + return x.left ?? x.right; + } + // find node that do not have a left child + // in the subtree of the given node + function successor(x: RBTreeNode): RBTreeNode { + let temp = x; + while (temp.left) temp = temp.left; + return temp; + } + } + + fixDoubleBlack(x: RBTreeNode): void { + if (x === this.root) return; // Reached root + + const sibling = x.sibling(); + const parent = x.parent!; + if (!sibling) { + // No sibiling, double black pushed up + this.fixDoubleBlack(parent); + } else { + if (sibling.color === 0) { + // Sibling red + parent.color = 0; + sibling.color = 1; + if (sibling.isOnLeft()) this.rotateRight(parent); + // left case + else this.rotateLeft(parent); // right case + this.fixDoubleBlack(x); + } else { + // Sibling black + if (sibling.hasRedChild()) { + // at least 1 red children + if (sibling.left && sibling.left.color === 0) { + if (sibling.isOnLeft()) { + // left left + sibling.left.color = sibling.color; + sibling.color = parent.color; + this.rotateRight(parent); + } else { + // right left + sibling.left.color = parent.color; + this.rotateRight(sibling); + this.rotateLeft(parent); + } + } else { + if (sibling.isOnLeft()) { + // left right + sibling.right!.color = parent.color; + this.rotateLeft(sibling); + this.rotateRight(parent); + } else { + // right right + sibling.right!.color = sibling.color; + sibling.color = parent.color; + this.rotateLeft(parent); + } + } + parent.color = 1; + } else { + // 2 black children + sibling.color = 0; + if (parent.color === 1) this.fixDoubleBlack(parent); + else parent.color = 1; + } + } + } + } + + insert(data: T): boolean { + // search for a position to insert + let parent = this.root; + while (parent) { + if (this.lt(data, parent.data)) { + if (!parent.left) break; + else parent = parent.left; + } else if (this.lt(parent.data, data)) { + if (!parent.right) break; + else parent = parent.right; + } else break; + } + + // insert node into parent + const node = new RBTreeNode(data); + if (!parent) this.root = node; + else if (this.lt(node.data, parent.data)) parent.left = node; + else if (this.lt(parent.data, node.data)) parent.right = node; + else { + parent.count++; + return false; + } + node.parent = parent; + this.fixAfterInsert(node); + return true; + } + + find(data: T): RBTreeNode | null { + let p = this.root; + while (p) { + if (this.lt(data, p.data)) { + p = p.left; + } else if (this.lt(p.data, data)) { + p = p.right; + } else break; + } + return p ?? null; + } + + *inOrder(root: RBTreeNode = this.root!): Generator { + if (!root) return; + for (const v of this.inOrder(root.left!)) yield v; + yield root.data; + for (const v of this.inOrder(root.right!)) yield v; + } + + *reverseInOrder(root: RBTreeNode = this.root!): Generator { + if (!root) return; + for (const v of this.reverseInOrder(root.right!)) yield v; + yield root.data; + for (const v of this.reverseInOrder(root.left!)) yield v; + } +} + +class TreeSet { + _size: number; + tree: RBTree; + compare: Compare; + constructor( + collection: T[] | Compare = [], + compare: Compare = (l: T, r: T) => (l < r ? -1 : l > r ? 1 : 0), + ) { + if (typeof collection === 'function') { + compare = collection; + collection = []; + } + this._size = 0; + this.compare = compare; + this.tree = new RBTree(compare); + for (const val of collection) this.add(val); + } + + size(): number { + return this._size; + } + + has(val: T): boolean { + return !!this.tree.find(val); + } + + add(val: T): boolean { + const successful = this.tree.insert(val); + this._size += successful ? 1 : 0; + return successful; + } + + delete(val: T): boolean { + const deleted = this.tree.deleteAll(val); + this._size -= deleted ? 1 : 0; + return deleted; + } + + ceil(val: T): T | undefined { + let p = this.tree.root; + let higher = null; + while (p) { + if (this.compare(p.data, val) >= 0) { + higher = p; + p = p.left; + } else { + p = p.right; + } + } + return higher?.data; + } + + floor(val: T): T | undefined { + let p = this.tree.root; + let lower = null; + while (p) { + if (this.compare(val, p.data) >= 0) { + lower = p; + p = p.right; + } else { + p = p.left; + } + } + return lower?.data; + } + + higher(val: T): T | undefined { + let p = this.tree.root; + let higher = null; + while (p) { + if (this.compare(val, p.data) < 0) { + higher = p; + p = p.left; + } else { + p = p.right; + } + } + return higher?.data; + } + + lower(val: T): T | undefined { + let p = this.tree.root; + let lower = null; + while (p) { + if (this.compare(p.data, val) < 0) { + lower = p; + p = p.right; + } else { + p = p.left; + } + } + return lower?.data; + } + + first(): T | undefined { + return this.tree.inOrder().next().value; + } + last(): T | undefined { + return this.tree.reverseInOrder().next().value; + } + + shift(): T | undefined { + const first = this.first(); + if (first === undefined) return undefined; + this.delete(first); + return first; + } + + pop(): T | undefined { + const last = this.last(); + if (last === undefined) return undefined; + this.delete(last); + return last; + } + + *[Symbol.iterator](): Generator { + for (const val of this.values()) yield val; + } + + *keys(): Generator { + for (const val of this.values()) yield val; + } + + *values(): Generator { + for (const val of this.tree.inOrder()) yield val; + return undefined; + } + + /** + * Return a generator for reverse order traversing the set + */ + *rvalues(): Generator { + for (const val of this.tree.reverseInOrder()) yield val; + return undefined; + } +} + +class TreeMultiSet { + _size: number; + tree: RBTree; + compare: Compare; + constructor( + collection: T[] | Compare = [], + compare: Compare = (l: T, r: T) => (l < r ? -1 : l > r ? 1 : 0), + ) { + if (typeof collection === 'function') { + compare = collection; + collection = []; + } + this._size = 0; + this.compare = compare; + this.tree = new RBTree(compare); + for (const val of collection) this.add(val); + } + + size(): number { + return this._size; + } + + has(val: T): boolean { + return !!this.tree.find(val); + } + + add(val: T): boolean { + const successful = this.tree.insert(val); + this._size++; + return successful; + } + + delete(val: T): boolean { + const successful = this.tree.delete(val); + if (!successful) return false; + this._size--; + return true; + } + + count(val: T): number { + const node = this.tree.find(val); + return node ? node.count : 0; + } + + ceil(val: T): T | undefined { + let p = this.tree.root; + let higher = null; + while (p) { + if (this.compare(p.data, val) >= 0) { + higher = p; + p = p.left; + } else { + p = p.right; + } + } + return higher?.data; + } + + floor(val: T): T | undefined { + let p = this.tree.root; + let lower = null; + while (p) { + if (this.compare(val, p.data) >= 0) { + lower = p; + p = p.right; + } else { + p = p.left; + } + } + return lower?.data; + } + + higher(val: T): T | undefined { + let p = this.tree.root; + let higher = null; + while (p) { + if (this.compare(val, p.data) < 0) { + higher = p; + p = p.left; + } else { + p = p.right; + } + } + return higher?.data; + } + + lower(val: T): T | undefined { + let p = this.tree.root; + let lower = null; + while (p) { + if (this.compare(p.data, val) < 0) { + lower = p; + p = p.right; + } else { + p = p.left; + } + } + return lower?.data; + } + + first(): T | undefined { + return this.tree.inOrder().next().value; + } + + last(): T | undefined { + return this.tree.reverseInOrder().next().value; + } + + shift(): T | undefined { + const first = this.first(); + if (first === undefined) return undefined; + this.delete(first); + return first; + } + + pop(): T | undefined { + const last = this.last(); + if (last === undefined) return undefined; + this.delete(last); + return last; + } + + *[Symbol.iterator](): Generator { + yield* this.values(); + } + + *keys(): Generator { + for (const val of this.values()) yield val; + } + + *values(): Generator { + for (const val of this.tree.inOrder()) { + let count = this.count(val); + while (count--) yield val; + } + return undefined; + } + + /** + * Return a generator for reverse order traversing the multi-set + */ + *rvalues(): Generator { + for (const val of this.tree.reverseInOrder()) { + let count = this.count(val); + while (count--) yield val; + } + return undefined; + } +} ``` ### **...** diff --git a/solution/2400-2499/2454.Next Greater Element IV/Solution.cpp b/solution/2400-2499/2454.Next Greater Element IV/Solution.cpp index 57d91102cb7fa..1f6c24bb54184 100644 --- a/solution/2400-2499/2454.Next Greater Element IV/Solution.cpp +++ b/solution/2400-2499/2454.Next Greater Element IV/Solution.cpp @@ -1,24 +1,21 @@ -using pii = pair; - -class Solution { -public: - vector secondGreaterElement(vector& nums) { - stack stk; - priority_queue, greater> q; - int n = nums.size(); - vector ans(n, -1); - for (int i = 0; i < n; ++i) { - int v = nums[i]; - while (!q.empty() && q.top().first < v) { - ans[q.top().second] = v; - q.pop(); - } - while (!stk.empty() && nums[stk.top()] < v) { - q.push({nums[stk.top()], stk.top()}); - stk.pop(); - } - stk.push(i); - } - return ans; - } -}; \ No newline at end of file +class Solution { +public: + vector secondGreaterElement(vector& nums) { + int n = nums.size(); + vector ans(n, -1); + vector> arr(n); + for (int i = 0; i < n; ++i) { + arr[i] = {-nums[i], i}; + } + sort(arr.begin(), arr.end()); + set ts; + for (auto& [_, i] : arr) { + auto it = ts.upper_bound(i); + if (it != ts.end() && ts.upper_bound(*it) != ts.end()) { + ans[i] = nums[*ts.upper_bound(*it)]; + } + ts.insert(i); + } + return ans; + } +}; diff --git a/solution/2400-2499/2454.Next Greater Element IV/Solution.go b/solution/2400-2499/2454.Next Greater Element IV/Solution.go deleted file mode 100644 index dd1aeaa772cb8..0000000000000 --- a/solution/2400-2499/2454.Next Greater Element IV/Solution.go +++ /dev/null @@ -1,34 +0,0 @@ -func secondGreaterElement(nums []int) []int { - stk := []int{} - q := hp{} - n := len(nums) - ans := make([]int, n) - for i := range ans { - ans[i] = -1 - } - for i, v := range nums { - for len(q) > 0 && q[0].v < v { - ans[q[0].i] = v - heap.Pop(&q) - } - for len(stk) > 0 && nums[stk[len(stk)-1]] < v { - heap.Push(&q, pair{nums[stk[len(stk)-1]], stk[len(stk)-1]}) - stk = stk[:len(stk)-1] - } - stk = append(stk, i) - } - return ans -} - -type pair struct{ v, i int } - -type hp []pair - -func (h hp) Len() int { return len(h) } -func (h hp) Less(i, j int) bool { - a, b := h[i], h[j] - return a.v < b.v -} -func (h hp) Swap(i, j int) { h[i], h[j] = h[j], h[i] } -func (h *hp) Push(v any) { *h = append(*h, v.(pair)) } -func (h *hp) Pop() any { a := *h; v := a[len(a)-1]; *h = a[:len(a)-1]; return v } \ No newline at end of file diff --git a/solution/2400-2499/2454.Next Greater Element IV/Solution.java b/solution/2400-2499/2454.Next Greater Element IV/Solution.java index 110c8f0165eb0..fcfa3bdb6ef18 100644 --- a/solution/2400-2499/2454.Next Greater Element IV/Solution.java +++ b/solution/2400-2499/2454.Next Greater Element IV/Solution.java @@ -1,21 +1,22 @@ -class Solution { - public int[] secondGreaterElement(int[] nums) { - Deque stk = new ArrayDeque<>(); - PriorityQueue q = new PriorityQueue<>((a, b) -> a[0] - b[0]); - int n = nums.length; - int[] ans = new int[n]; - Arrays.fill(ans, -1); - for (int i = 0; i < n; ++i) { - int v = nums[i]; - while (!q.isEmpty() && q.peek()[0] < v) { - ans[q.peek()[1]] = v; - q.poll(); - } - while (!stk.isEmpty() && nums[stk.peek()] < v) { - q.offer(new int[] {nums[stk.peek()], stk.pop()}); - } - stk.push(i); - } - return ans; - } +class Solution { + public int[] secondGreaterElement(int[] nums) { + int n = nums.length; + int[] ans = new int[n]; + Arrays.fill(ans, -1); + int[][] arr = new int[n][0]; + for (int i = 0; i < n; ++i) { + arr[i] = new int[] {nums[i], i}; + } + Arrays.sort(arr, (a, b) -> b[0] - a[0]); + TreeSet ts = new TreeSet<>(); + for (int[] pair : arr) { + int i = pair[1]; + Integer j = ts.higher(i); + if (j != null && ts.higher(j) != null) { + ans[i] = nums[ts.higher(j)]; + } + ts.add(i); + } + return ans; + } } \ No newline at end of file diff --git a/solution/2400-2499/2454.Next Greater Element IV/Solution.py b/solution/2400-2499/2454.Next Greater Element IV/Solution.py index 75b58be1dc698..8596b7950b61f 100644 --- a/solution/2400-2499/2454.Next Greater Element IV/Solution.py +++ b/solution/2400-2499/2454.Next Greater Element IV/Solution.py @@ -1,13 +1,16 @@ -class Solution: - def secondGreaterElement(self, nums: List[int]) -> List[int]: - stk = [] - q = [] - ans = [-1] * len(nums) - for i, v in enumerate(nums): - while q and q[0][0] < v: - ans[q[0][1]] = v - heappop(q) - while stk and nums[stk[-1]] < v: - heappush(q, (nums[stk[-1]], stk.pop())) - stk.append(i) - return ans +from sortedcontainers import SortedList + + +class Solution: + def secondGreaterElement(self, nums: List[int]) -> List[int]: + arr = [(x, i) for i, x in enumerate(nums)] + arr.sort(key=lambda x: -x[0]) + sl = SortedList() + n = len(nums) + ans = [-1] * n + for _, i in arr: + j = sl.bisect_right(i) + if j + 1 < len(sl): + ans[i] = nums[sl[j + 1]] + sl.add(i) + return ans diff --git a/solution/2400-2499/2454.Next Greater Element IV/Solution.ts b/solution/2400-2499/2454.Next Greater Element IV/Solution.ts new file mode 100644 index 0000000000000..ca95a6a4e7bcf --- /dev/null +++ b/solution/2400-2499/2454.Next Greater Element IV/Solution.ts @@ -0,0 +1,661 @@ +function secondGreaterElement(nums: number[]): number[] { + const n = nums.length; + const arr: number[][] = []; + for (let i = 0; i < n; ++i) { + arr.push([nums[i], i]); + } + arr.sort((a, b) => b[0] - a[0]); + const ans = Array(n).fill(-1); + const ts = new TreeSet(); + for (const [_, i] of arr) { + let j = ts.higher(i); + if (j !== undefined) { + j = ts.higher(j); + if (j !== undefined) { + ans[i] = nums[j]; + } + } + ts.add(i); + } + return ans; +} + +type Compare = (lhs: T, rhs: T) => number; + +class RBTreeNode { + data: T; + count: number; + left: RBTreeNode | null; + right: RBTreeNode | null; + parent: RBTreeNode | null; + color: number; + constructor(data: T) { + this.data = data; + this.left = this.right = this.parent = null; + this.color = 0; + this.count = 1; + } + + sibling(): RBTreeNode | null { + if (!this.parent) return null; // sibling null if no parent + return this.isOnLeft() ? this.parent.right : this.parent.left; + } + + isOnLeft(): boolean { + return this === this.parent!.left; + } + + hasRedChild(): boolean { + return ( + Boolean(this.left && this.left.color === 0) || + Boolean(this.right && this.right.color === 0) + ); + } +} + +class RBTree { + root: RBTreeNode | null; + lt: (l: T, r: T) => boolean; + constructor(compare: Compare = (l: T, r: T) => (l < r ? -1 : l > r ? 1 : 0)) { + this.root = null; + this.lt = (l: T, r: T) => compare(l, r) < 0; + } + + rotateLeft(pt: RBTreeNode): void { + const right = pt.right!; + pt.right = right.left; + + if (pt.right) pt.right.parent = pt; + right.parent = pt.parent; + + if (!pt.parent) this.root = right; + else if (pt === pt.parent.left) pt.parent.left = right; + else pt.parent.right = right; + + right.left = pt; + pt.parent = right; + } + + rotateRight(pt: RBTreeNode): void { + const left = pt.left!; + pt.left = left.right; + + if (pt.left) pt.left.parent = pt; + left.parent = pt.parent; + + if (!pt.parent) this.root = left; + else if (pt === pt.parent.left) pt.parent.left = left; + else pt.parent.right = left; + + left.right = pt; + pt.parent = left; + } + + swapColor(p1: RBTreeNode, p2: RBTreeNode): void { + const tmp = p1.color; + p1.color = p2.color; + p2.color = tmp; + } + + swapData(p1: RBTreeNode, p2: RBTreeNode): void { + const tmp = p1.data; + p1.data = p2.data; + p2.data = tmp; + } + + fixAfterInsert(pt: RBTreeNode): void { + let parent = null; + let grandParent = null; + + while (pt !== this.root && pt.color !== 1 && pt.parent?.color === 0) { + parent = pt.parent; + grandParent = pt.parent.parent; + + /* Case : A + Parent of pt is left child of Grand-parent of pt */ + if (parent === grandParent?.left) { + const uncle = grandParent.right; + + /* Case : 1 + The uncle of pt is also red + Only Recoloring required */ + if (uncle && uncle.color === 0) { + grandParent.color = 0; + parent.color = 1; + uncle.color = 1; + pt = grandParent; + } else { + /* Case : 2 + pt is right child of its parent + Left-rotation required */ + if (pt === parent.right) { + this.rotateLeft(parent); + pt = parent; + parent = pt.parent; + } + + /* Case : 3 + pt is left child of its parent + Right-rotation required */ + this.rotateRight(grandParent); + this.swapColor(parent!, grandParent); + pt = parent!; + } + } else { + /* Case : B + Parent of pt is right child of Grand-parent of pt */ + const uncle = grandParent!.left; + + /* Case : 1 + The uncle of pt is also red + Only Recoloring required */ + if (uncle != null && uncle.color === 0) { + grandParent!.color = 0; + parent.color = 1; + uncle.color = 1; + pt = grandParent!; + } else { + /* Case : 2 + pt is left child of its parent + Right-rotation required */ + if (pt === parent.left) { + this.rotateRight(parent); + pt = parent; + parent = pt.parent; + } + + /* Case : 3 + pt is right child of its parent + Left-rotation required */ + this.rotateLeft(grandParent!); + this.swapColor(parent!, grandParent!); + pt = parent!; + } + } + } + this.root!.color = 1; + } + + delete(val: T): boolean { + const node = this.find(val); + if (!node) return false; + node.count--; + if (!node.count) this.deleteNode(node); + return true; + } + + deleteAll(val: T): boolean { + const node = this.find(val); + if (!node) return false; + this.deleteNode(node); + return true; + } + + deleteNode(v: RBTreeNode): void { + const u = BSTreplace(v); + + // True when u and v are both black + const uvBlack = (u === null || u.color === 1) && v.color === 1; + const parent = v.parent!; + + if (!u) { + // u is null therefore v is leaf + if (v === this.root) this.root = null; + // v is root, making root null + else { + if (uvBlack) { + // u and v both black + // v is leaf, fix double black at v + this.fixDoubleBlack(v); + } else { + // u or v is red + if (v.sibling()) { + // sibling is not null, make it red" + v.sibling()!.color = 0; + } + } + // delete v from the tree + if (v.isOnLeft()) parent.left = null; + else parent.right = null; + } + return; + } + + if (!v.left || !v.right) { + // v has 1 child + if (v === this.root) { + // v is root, assign the value of u to v, and delete u + v.data = u.data; + v.left = v.right = null; + } else { + // Detach v from tree and move u up + if (v.isOnLeft()) parent.left = u; + else parent.right = u; + u.parent = parent; + if (uvBlack) this.fixDoubleBlack(u); + // u and v both black, fix double black at u + else u.color = 1; // u or v red, color u black + } + return; + } + + // v has 2 children, swap data with successor and recurse + this.swapData(u, v); + this.deleteNode(u); + + // find node that replaces a deleted node in BST + function BSTreplace(x: RBTreeNode): RBTreeNode | null { + // when node have 2 children + if (x.left && x.right) return successor(x.right); + // when leaf + if (!x.left && !x.right) return null; + // when single child + return x.left ?? x.right; + } + // find node that do not have a left child + // in the subtree of the given node + function successor(x: RBTreeNode): RBTreeNode { + let temp = x; + while (temp.left) temp = temp.left; + return temp; + } + } + + fixDoubleBlack(x: RBTreeNode): void { + if (x === this.root) return; // Reached root + + const sibling = x.sibling(); + const parent = x.parent!; + if (!sibling) { + // No sibiling, double black pushed up + this.fixDoubleBlack(parent); + } else { + if (sibling.color === 0) { + // Sibling red + parent.color = 0; + sibling.color = 1; + if (sibling.isOnLeft()) this.rotateRight(parent); + // left case + else this.rotateLeft(parent); // right case + this.fixDoubleBlack(x); + } else { + // Sibling black + if (sibling.hasRedChild()) { + // at least 1 red children + if (sibling.left && sibling.left.color === 0) { + if (sibling.isOnLeft()) { + // left left + sibling.left.color = sibling.color; + sibling.color = parent.color; + this.rotateRight(parent); + } else { + // right left + sibling.left.color = parent.color; + this.rotateRight(sibling); + this.rotateLeft(parent); + } + } else { + if (sibling.isOnLeft()) { + // left right + sibling.right!.color = parent.color; + this.rotateLeft(sibling); + this.rotateRight(parent); + } else { + // right right + sibling.right!.color = sibling.color; + sibling.color = parent.color; + this.rotateLeft(parent); + } + } + parent.color = 1; + } else { + // 2 black children + sibling.color = 0; + if (parent.color === 1) this.fixDoubleBlack(parent); + else parent.color = 1; + } + } + } + } + + insert(data: T): boolean { + // search for a position to insert + let parent = this.root; + while (parent) { + if (this.lt(data, parent.data)) { + if (!parent.left) break; + else parent = parent.left; + } else if (this.lt(parent.data, data)) { + if (!parent.right) break; + else parent = parent.right; + } else break; + } + + // insert node into parent + const node = new RBTreeNode(data); + if (!parent) this.root = node; + else if (this.lt(node.data, parent.data)) parent.left = node; + else if (this.lt(parent.data, node.data)) parent.right = node; + else { + parent.count++; + return false; + } + node.parent = parent; + this.fixAfterInsert(node); + return true; + } + + find(data: T): RBTreeNode | null { + let p = this.root; + while (p) { + if (this.lt(data, p.data)) { + p = p.left; + } else if (this.lt(p.data, data)) { + p = p.right; + } else break; + } + return p ?? null; + } + + *inOrder(root: RBTreeNode = this.root!): Generator { + if (!root) return; + for (const v of this.inOrder(root.left!)) yield v; + yield root.data; + for (const v of this.inOrder(root.right!)) yield v; + } + + *reverseInOrder(root: RBTreeNode = this.root!): Generator { + if (!root) return; + for (const v of this.reverseInOrder(root.right!)) yield v; + yield root.data; + for (const v of this.reverseInOrder(root.left!)) yield v; + } +} + +class TreeSet { + _size: number; + tree: RBTree; + compare: Compare; + constructor( + collection: T[] | Compare = [], + compare: Compare = (l: T, r: T) => (l < r ? -1 : l > r ? 1 : 0), + ) { + if (typeof collection === 'function') { + compare = collection; + collection = []; + } + this._size = 0; + this.compare = compare; + this.tree = new RBTree(compare); + for (const val of collection) this.add(val); + } + + size(): number { + return this._size; + } + + has(val: T): boolean { + return !!this.tree.find(val); + } + + add(val: T): boolean { + const successful = this.tree.insert(val); + this._size += successful ? 1 : 0; + return successful; + } + + delete(val: T): boolean { + const deleted = this.tree.deleteAll(val); + this._size -= deleted ? 1 : 0; + return deleted; + } + + ceil(val: T): T | undefined { + let p = this.tree.root; + let higher = null; + while (p) { + if (this.compare(p.data, val) >= 0) { + higher = p; + p = p.left; + } else { + p = p.right; + } + } + return higher?.data; + } + + floor(val: T): T | undefined { + let p = this.tree.root; + let lower = null; + while (p) { + if (this.compare(val, p.data) >= 0) { + lower = p; + p = p.right; + } else { + p = p.left; + } + } + return lower?.data; + } + + higher(val: T): T | undefined { + let p = this.tree.root; + let higher = null; + while (p) { + if (this.compare(val, p.data) < 0) { + higher = p; + p = p.left; + } else { + p = p.right; + } + } + return higher?.data; + } + + lower(val: T): T | undefined { + let p = this.tree.root; + let lower = null; + while (p) { + if (this.compare(p.data, val) < 0) { + lower = p; + p = p.right; + } else { + p = p.left; + } + } + return lower?.data; + } + + first(): T | undefined { + return this.tree.inOrder().next().value; + } + + last(): T | undefined { + return this.tree.reverseInOrder().next().value; + } + + shift(): T | undefined { + const first = this.first(); + if (first === undefined) return undefined; + this.delete(first); + return first; + } + + pop(): T | undefined { + const last = this.last(); + if (last === undefined) return undefined; + this.delete(last); + return last; + } + + *[Symbol.iterator](): Generator { + for (const val of this.values()) yield val; + } + + *keys(): Generator { + for (const val of this.values()) yield val; + } + + *values(): Generator { + for (const val of this.tree.inOrder()) yield val; + return undefined; + } + + /** + * Return a generator for reverse order traversing the set + */ + *rvalues(): Generator { + for (const val of this.tree.reverseInOrder()) yield val; + return undefined; + } +} + +class TreeMultiSet { + _size: number; + tree: RBTree; + compare: Compare; + constructor( + collection: T[] | Compare = [], + compare: Compare = (l: T, r: T) => (l < r ? -1 : l > r ? 1 : 0), + ) { + if (typeof collection === 'function') { + compare = collection; + collection = []; + } + this._size = 0; + this.compare = compare; + this.tree = new RBTree(compare); + for (const val of collection) this.add(val); + } + + size(): number { + return this._size; + } + + has(val: T): boolean { + return !!this.tree.find(val); + } + + add(val: T): boolean { + const successful = this.tree.insert(val); + this._size++; + return successful; + } + + delete(val: T): boolean { + const successful = this.tree.delete(val); + if (!successful) return false; + this._size--; + return true; + } + + count(val: T): number { + const node = this.tree.find(val); + return node ? node.count : 0; + } + + ceil(val: T): T | undefined { + let p = this.tree.root; + let higher = null; + while (p) { + if (this.compare(p.data, val) >= 0) { + higher = p; + p = p.left; + } else { + p = p.right; + } + } + return higher?.data; + } + + floor(val: T): T | undefined { + let p = this.tree.root; + let lower = null; + while (p) { + if (this.compare(val, p.data) >= 0) { + lower = p; + p = p.right; + } else { + p = p.left; + } + } + return lower?.data; + } + + higher(val: T): T | undefined { + let p = this.tree.root; + let higher = null; + while (p) { + if (this.compare(val, p.data) < 0) { + higher = p; + p = p.left; + } else { + p = p.right; + } + } + return higher?.data; + } + + lower(val: T): T | undefined { + let p = this.tree.root; + let lower = null; + while (p) { + if (this.compare(p.data, val) < 0) { + lower = p; + p = p.right; + } else { + p = p.left; + } + } + return lower?.data; + } + + first(): T | undefined { + return this.tree.inOrder().next().value; + } + + last(): T | undefined { + return this.tree.reverseInOrder().next().value; + } + + shift(): T | undefined { + const first = this.first(); + if (first === undefined) return undefined; + this.delete(first); + return first; + } + + pop(): T | undefined { + const last = this.last(); + if (last === undefined) return undefined; + this.delete(last); + return last; + } + + *[Symbol.iterator](): Generator { + yield* this.values(); + } + + *keys(): Generator { + for (const val of this.values()) yield val; + } + + *values(): Generator { + for (const val of this.tree.inOrder()) { + let count = this.count(val); + while (count--) yield val; + } + return undefined; + } + + /** + * Return a generator for reverse order traversing the multi-set + */ + *rvalues(): Generator { + for (const val of this.tree.reverseInOrder()) { + let count = this.count(val); + while (count--) yield val; + } + return undefined; + } +}