Welcome to Subscribe On Youtube

Formatted question description: https://leetcode.ca/all/480.html

480. Sliding Window Median

Level

Hard

Description

Median is the middle value in an ordered integer list. If the size of the list is even, there is no middle value. So the median is the mean of the two middle value.

Examples:

[2,3,4], the median is 3

[2,3], the median is (2 + 3) / 2 = 2.5

Given an array nums, there is a sliding window of size k which is moving from the very left of the array to the very right. You can only see the k numbers in the window. Each time the sliding window moves right by one position. Your job is to output the median array for each window in the original array.

For example,

Given nums = [1,3,-1,-3,5,3,6,7], and k = 3.

Window position                Median
---------------               -----
[1  3  -1] -3  5  3  6  7       1
 1 [3  -1  -3] 5  3  6  7       -1
 1  3 [-1  -3  5] 3  6  7       -1
 1  3  -1 [-3  5  3] 6  7       3
 1  3  -1  -3 [5  3  6] 7       5
 1  3  -1  -3  5 [3  6  7]      6

Therefore, return the median sliding window as [1,-1,-1,3,5,6].

Note:

You may assume k is always valid, ie: k is always smaller than input array’s size for non-empty array.

Answers within 10^-5 of the actual value will be accepted as correct.

Solution

Solution 1 - Two heaps

Use two priority queues to store the smaller half elements and the larger half elements respectively. The first priority queue priorityQueue1 stores the smaller half elements and the largest element is polled first. The second priority queue priorityQueue2 stores the larger half elements and the smallest element is polled first. If there are odd number of elements, then priorityQueue1 stores one more element than priorityQueue2.

The median is calculated in a straightforward way. If k is odd, then peek one element from priorityQueue1, which is the median. If k is even, then peek one element from each queue priorityQueue1 and priorityQueue2 respectively and calculate the mean of the two peeked elements, which is the median.

Since this problem requires calculating the median array for each window, polling elements and offering elements frequently is not efficient. A better way is to use a map to store the invalid elements and the number of occurrences of each invalid element. Each time obtain the previous element in the window that is not in the window any longer and update its number of occurrences in the map. After the previous element and the new element are obtained, check which priority queues the two elements should be in, comparing the two elements with the elements at the top of both priority queues. Then add the new element to the correct priority queue. Then for both priority queues, while the element at the top of each priority queue is in the map of invalid elements, poll them out and update the map. Then calculate the median accordingly.

Solution 2 - Sorted list-k

We first add the first k arrays of the array to the set. Since multiset has its own sorting function, we can quickly find the iterator mid pointing to the middle number through k/2.

  • If k is an odd number, then the number pointed to by mid which is the median;
  • if k is an even number, then the average of the number pointed to by mid and the previous number is the median.

When we add a new number to the collection, the multiset will be added to the correct position according to the size of the new number, and then we see that

  • if the newly added number is smaller than the number pointed to by the previous mid, then the median must be pulled down , so mid moves forward by one,
  • and if the number to be deleted is less than or equal to the number pointed to by mid (note that the equal sign is added here because the number to be deleted may be the number pointed to by mid), then mid moves backward by one.

Then we delete the leftmost number of the sliding window.

  • class MedianFinder {
        private PriorityQueue<Integer> small = new PriorityQueue<>(Comparator.reverseOrder());
        private PriorityQueue<Integer> large = new PriorityQueue<>();
        private Map<Integer, Integer> delayed = new HashMap<>();
        private int smallSize;
        private int largeSize;
        private int k;
    
        public MedianFinder(int k) {
            this.k = k;
        }
    
        public void addNum(int num) {
            if (small.isEmpty() || num <= small.peek()) {
                small.offer(num);
                ++smallSize;
            } else {
                large.offer(num);
                ++largeSize;
            }
            rebalance();
        }
    
        public double findMedian() {
            return (k & 1) == 1 ? small.peek() : ((double) small.peek() + large.peek()) / 2;
        }
    
        public void removeNum(int num) {
            delayed.merge(num, 1, Integer::sum);
            if (num <= small.peek()) {
                --smallSize;
                if (num == small.peek()) {
                    prune(small);
                }
            } else {
                --largeSize;
                if (num == large.peek()) {
                    prune(large);
                }
            }
            rebalance();
        }
    
        private void prune(PriorityQueue<Integer> pq) {
            while (!pq.isEmpty() && delayed.containsKey(pq.peek())) {
                if (delayed.merge(pq.peek(), -1, Integer::sum) == 0) {
                    delayed.remove(pq.peek());
                }
                pq.poll();
            }
        }
    
        private void rebalance() {
            if (smallSize > largeSize + 1) {
                large.offer(small.poll());
                --smallSize;
                ++largeSize;
                prune(small);
            } else if (smallSize < largeSize) {
                small.offer(large.poll());
                --largeSize;
                ++smallSize;
                prune(large);
            }
        }
    }
    
    
    class Solution {
        public double[] medianSlidingWindow(int[] nums, int k) {
            MedianFinder finder = new MedianFinder(k);
            for (int i = 0; i < k; ++i) {
                finder.addNum(nums[i]);
            }
            int n = nums.length;
            double[] ans = new double[n - k + 1];
            ans[0] = finder.findMedian();
            for (int i = k; i < n; ++i) {
                finder.addNum(nums[i]);
                finder.removeNum(nums[i - k]);
                ans[i - k + 1] = finder.findMedian();
            }
            return ans;
        }
    }
    
  • class Solution {
    public:
        vector<double> medianSlidingWindow(vector<int>& nums, int k) {
            vector<double> res;
            multiset<int> small, large;
            for (int i = 0; i < nums.size(); ++i) {
                if (i >= k) {
                    if (small.count(nums[i - k])) small.erase(small.find(nums[i - k]));
                    else if (large.count(nums[i - k])) large.erase(large.find(nums[i - k]));
                }
                if (small.size() <= large.size()) {
                    if (large.empty() || nums[i] <= *large.begin()) small.insert(nums[i]);
                    else {
                        small.insert(*large.begin());
                        large.erase(large.begin());
                        large.insert(nums[i]);
                    }
                } else {
                    if (nums[i] >= *small.rbegin()) large.insert(nums[i]);
                    else {
                        large.insert(*small.rbegin());
                        small.erase(--small.end());
                        small.insert(nums[i]);
                    }
                }
                if (i >= (k - 1)) {
                    if (k % 2) res.push_back(*small.rbegin());
                    else res.push_back(((double)*small.rbegin() + *large.begin()) / 2);
                }
            }
            return res;
        }
    };
    
  • class MedianFinder:
        def __init__(self, k: int):
            self.k = k
            self.small = []
            self.large = []
            self.delayed = defaultdict(int)
            self.small_size = 0
            self.large_size = 0
    
        def add_num(self, num: int):
            if not self.small or num <= -self.small[0]:
                heappush(self.small, -num)
                self.small_size += 1
            else:
                heappush(self.large, num)
                self.large_size += 1
            self.rebalance()
    
        def find_median(self) -> float:
            return -self.small[0] if self.k & 1 else (-self.small[0] + self.large[0]) / 2
    
        def remove_num(self, num: int):
            self.delayed[num] += 1
            if num <= -self.small[0]:
                self.small_size -= 1
                if num == -self.small[0]:
                    self.prune(self.small)
            else:
                self.large_size -= 1
                if num == self.large[0]:
                    self.prune(self.large)
            self.rebalance()
    
        def prune(self, pq: List[int]):
            sign = -1 if pq is self.small else 1
            while pq and sign * pq[0] in self.delayed:
                self.delayed[sign * pq[0]] -= 1
                if self.delayed[sign * pq[0]] == 0:
                    self.delayed.pop(sign * pq[0])
                heappop(pq)
    
        def rebalance(self):
            if self.small_size > self.large_size + 1:
                heappush(self.large, -heappop(self.small))
                self.small_size -= 1
                self.large_size += 1
                self.prune(self.small)
            elif self.small_size < self.large_size:
                heappush(self.small, -heappop(self.large))
                self.large_size -= 1
                self.small_size += 1
                self.prune(self.large)
    
    
    class Solution:
        def medianSlidingWindow(self, nums: List[int], k: int) -> List[float]:
            finder = MedianFinder(k)
            for x in nums[:k]:
                finder.add_num(x)
            ans = [finder.find_median()]
            for i in range(k, len(nums)):
                finder.add_num(nums[i])
                finder.remove_num(nums[i - k])
                ans.append(finder.find_median())
            return ans
    
    
  • type MedianFinder struct {
    	small                hp
    	large                hp
    	delayed              map[int]int
    	smallSize, largeSize int
    	k                    int
    }
    
    func Constructor(k int) MedianFinder {
    	return MedianFinder{hp{}, hp{}, map[int]int{}, 0, 0, k}
    }
    
    func (this *MedianFinder) AddNum(num int) {
    	if this.small.Len() == 0 || num <= -this.small.IntSlice[0] {
    		heap.Push(&this.small, -num)
    		this.smallSize++
    	} else {
    		heap.Push(&this.large, num)
    		this.largeSize++
    	}
    	this.rebalance()
    }
    
    func (this *MedianFinder) FindMedian() float64 {
    	if this.k&1 == 1 {
    		return float64(-this.small.IntSlice[0])
    	}
    	return float64(-this.small.IntSlice[0]+this.large.IntSlice[0]) / 2
    }
    
    func (this *MedianFinder) removeNum(num int) {
    	this.delayed[num]++
    	if num <= -this.small.IntSlice[0] {
    		this.smallSize--
    		if num == -this.small.IntSlice[0] {
    			this.prune(&this.small)
    		}
    	} else {
    		this.largeSize--
    		if num == this.large.IntSlice[0] {
    			this.prune(&this.large)
    		}
    	}
    	this.rebalance()
    }
    
    func (this *MedianFinder) prune(pq *hp) {
    	sign := 1
    	if pq == &this.small {
    		sign = -1
    	}
    	for pq.Len() > 0 && this.delayed[sign*pq.IntSlice[0]] > 0 {
    		this.delayed[sign*pq.IntSlice[0]]--
    		if this.delayed[sign*pq.IntSlice[0]] == 0 {
    			delete(this.delayed, sign*pq.IntSlice[0])
    		}
    		heap.Pop(pq)
    	}
    }
    
    func (this *MedianFinder) rebalance() {
    	if this.smallSize > this.largeSize+1 {
    		heap.Push(&this.large, -heap.Pop(&this.small).(int))
    		this.smallSize--
    		this.largeSize++
    		this.prune(&this.small)
    	} else if this.smallSize < this.largeSize {
    		heap.Push(&this.small, -heap.Pop(&this.large).(int))
    		this.smallSize++
    		this.largeSize--
    		this.prune(&this.large)
    	}
    }
    
    func medianSlidingWindow(nums []int, k int) []float64 {
    	finder := Constructor(k)
    	for _, num := range nums[:k] {
    		finder.AddNum(num)
    	}
    	ans := []float64{finder.FindMedian()}
    	for i := k; i < len(nums); i++ {
    		finder.AddNum(nums[i])
    		finder.removeNum(nums[i-k])
    		ans = append(ans, finder.FindMedian())
    	}
    	return ans
    }
    
    type hp struct{ sort.IntSlice }
    
    func (h hp) Less(i, j int) bool  { return h.IntSlice[i] < h.IntSlice[j] }
    func (h *hp) Push(v interface{}) { h.IntSlice = append(h.IntSlice, v.(int)) }
    func (h *hp) Pop() interface{} {
    	a := h.IntSlice
    	v := a[len(a)-1]
    	h.IntSlice = a[:len(a)-1]
    	return v
    }
    
  • class Solution {
        public double[] medianSlidingWindow(int[] nums, int k) {
            TreeMap<Integer, Integer> l = new TreeMap<>();
            TreeMap<Integer, Integer> r = new TreeMap<>();
            int n = nums.length;
            double[] ans = new double[n - k + 1];
            int lSize = 0, rSize = 0;
            for (int i = 0; i < n; ++i) {
                r.merge(nums[i], 1, Integer::sum);
                int x = r.firstKey();
                if (r.merge(x, -1, Integer::sum) == 0) {
                    r.remove(x);
                }
                l.merge(x, 1, Integer::sum);
                ++lSize;
                while (lSize - rSize > 1) {
                    x = l.lastKey();
                    if (l.merge(x, -1, Integer::sum) == 0) {
                        l.remove(x);
                    }
                    r.merge(x, 1, Integer::sum);
                    --lSize;
                    ++rSize;
                }
                int j = i - k + 1;
                if (j >= 0) {
                    ans[j] = k % 2 == 1 ? l.lastKey() : ((double) l.lastKey() + r.firstKey()) / 2;
                    if (l.containsKey(nums[j])) {
                        if (l.merge(nums[j], -1, Integer::sum) == 0) {
                            l.remove(nums[j]);
                        }
                        --lSize;
                    } else {
                        if (r.merge(nums[j], -1, Integer::sum) == 0) {
                            r.remove(nums[j]);
                        }
                        --rSize;
                    }
                }
            }
            return ans;
        }
    }
    
    
  • class Solution {
    public:
        vector<double> medianSlidingWindow(vector<int>& nums, int k) {
            multiset<int> l, r;
            int n = nums.size();
            vector<double> ans;
            for (int i = 0; i < n; ++i) {
                r.insert(nums[i]);
                l.insert(*r.begin());
                r.erase(r.begin());
                while (l.size() - r.size() > 1) {
                    r.insert(*l.rbegin());
                    l.erase(prev(l.end()));
                }
                int j = i - k + 1;
                if (j >= 0) {
                    ans.push_back(k % 2 ? *l.rbegin() : ((double) *l.rbegin() + *r.begin()) / 2);
                    auto it = l.find(nums[j]);
                    if (it != l.end()) {
                        l.erase(it);
                    } else {
                        r.erase(r.find(nums[j]));
                    }
                }
            }
            return ans;
        }
    };
    
    
  • class Solution:
        def medianSlidingWindow(self, nums: List[int], k: int) -> List[float]:
            l = SortedList()
            r = SortedList()
            ans = []
            for i, x in enumerate(nums):
                r.add(x)
                l.add(r.pop(0))
                while len(l) - len(r) > 1:
                    r.add(l.pop())
                j = i - k + 1
                if j >= 0:
                    ans.append(l[-1] if k & 1 else (l[-1] + r[0]) / 2)
                    if nums[j] in l:
                        l.remove(nums[j])
                    else:
                        r.remove(nums[j])
            return ans
    
    
  • func medianSlidingWindow(nums []int, k int) (ans []float64) {
    	l := redblacktree.New[int, int]()
    	r := redblacktree.New[int, int]()
    	merge := func(st *redblacktree.Tree[int, int], x, v int) {
    		c, _ := st.Get(x)
    		if c+v == 0 {
    			st.Remove(x)
    		} else {
    			st.Put(x, c+v)
    		}
    	}
    	lSize, rSize := 0, 0
    	for i, x := range nums {
    		merge(r, x, 1)
    		x = r.Left().Key
    		merge(r, x, -1)
    		merge(l, x, 1)
    		lSize++
    		for lSize-rSize > 1 {
    			x = l.Right().Key
    			merge(l, x, -1)
    			merge(r, x, 1)
    			lSize--
    			rSize++
    		}
    		if j := i - k + 1; j >= 0 {
    			if k%2 == 1 {
    				ans = append(ans, float64(l.Right().Key))
    			} else {
    				ans = append(ans, float64(l.Right().Key+r.Left().Key)/2)
    			}
    			if x = nums[j]; x <= l.Right().Key {
    				merge(l, x, -1)
    				lSize--
    			} else {
    				merge(r, x, -1)
    				rSize--
    			}
    		}
    	}
    	return
    }
    
    
  • function medianSlidingWindow(nums: number[], k: number): number[] {
        const l = new TreapMultiSet<number>((a, b) => a - b);
        const r = new TreapMultiSet<number>((a, b) => a - b);
        const n = nums.length;
        const ans: number[] = [];
        for (let i = 0; i < n; ++i) {
            r.add(nums[i]);
            l.add(r.shift()!);
            while (l.size - r.size > 1) {
                r.add(l.pop()!);
            }
            const j = i - k + 1;
            if (j >= 0) {
                ans[j] = k % 2 ? l.last()! : (l.last()! + r.first()!) / 2;
                if (nums[j] <= l.last()!) {
                    l.delete(nums[j]);
                } else {
                    r.delete(nums[j]);
                }
            }
        }
        return ans;
    }
    
    type CompareFunction<T, R extends 'number' | 'boolean'> = (
        a: T,
        b: T,
    ) => R extends 'number' ? number : boolean;
    
    interface ITreapMultiSet<T> extends Iterable<T> {
        add: (...value: T[]) => this;
        has: (value: T) => boolean;
        delete: (value: T) => void;
    
        bisectLeft: (value: T) => number;
        bisectRight: (value: T) => number;
    
        indexOf: (value: T) => number;
        lastIndexOf: (value: T) => number;
    
        at: (index: number) => T | undefined;
        first: () => T | undefined;
        last: () => T | undefined;
    
        lower: (value: T) => T | undefined;
        higher: (value: T) => T | undefined;
        floor: (value: T) => T | undefined;
        ceil: (value: T) => T | undefined;
    
        shift: () => T | undefined;
        pop: (index?: number) => T | undefined;
    
        count: (value: T) => number;
    
        keys: () => IterableIterator<T>;
        values: () => IterableIterator<T>;
        rvalues: () => IterableIterator<T>;
        entries: () => IterableIterator<[number, T]>;
    
        readonly size: number;
    }
    
    class TreapNode<T = number> {
        value: T;
        count: number;
        size: number;
        priority: number;
        left: TreapNode<T> | null;
        right: TreapNode<T> | null;
    
        constructor(value: T) {
            this.value = value;
            this.count = 1;
            this.size = 1;
            this.priority = Math.random();
            this.left = null;
            this.right = null;
        }
    
        static getSize(node: TreapNode<any> | null): number {
            return node?.size ?? 0;
        }
    
        static getFac(node: TreapNode<any> | null): number {
            return node?.priority ?? 0;
        }
    
        pushUp(): void {
            let tmp = this.count;
            tmp += TreapNode.getSize(this.left);
            tmp += TreapNode.getSize(this.right);
            this.size = tmp;
        }
    
        rotateRight(): TreapNode<T> {
            // eslint-disable-next-line @typescript-eslint/no-this-alias
            let node: TreapNode<T> = this;
            const left = node.left;
            node.left = left?.right ?? null;
            left && (left.right = node);
            left && (node = left);
            node.right?.pushUp();
            node.pushUp();
            return node;
        }
    
        rotateLeft(): TreapNode<T> {
            // eslint-disable-next-line @typescript-eslint/no-this-alias
            let node: TreapNode<T> = this;
            const right = node.right;
            node.right = right?.left ?? null;
            right && (right.left = node);
            right && (node = right);
            node.left?.pushUp();
            node.pushUp();
            return node;
        }
    }
    
    class TreapMultiSet<T = number> implements ITreapMultiSet<T> {
        private readonly root: TreapNode<T>;
        private readonly compareFn: CompareFunction<T, 'number'>;
        private readonly leftBound: T;
        private readonly rightBound: T;
    
        constructor(compareFn?: CompareFunction<T, 'number'>);
        constructor(compareFn: CompareFunction<T, 'number'>, leftBound: T, rightBound: T);
        constructor(
            compareFn: CompareFunction<T, any> = (a: any, b: any) => a - b,
            leftBound: any = -Infinity,
            rightBound: any = Infinity,
        ) {
            this.root = new TreapNode<T>(rightBound);
            this.root.priority = Infinity;
            this.root.left = new TreapNode<T>(leftBound);
            this.root.left.priority = -Infinity;
            this.root.pushUp();
    
            this.leftBound = leftBound;
            this.rightBound = rightBound;
            this.compareFn = compareFn;
        }
    
        get size(): number {
            return this.root.size - 2;
        }
    
        get height(): number {
            const getHeight = (node: TreapNode<T> | null): number => {
                if (node == null) return 0;
                return 1 + Math.max(getHeight(node.left), getHeight(node.right));
            };
    
            return getHeight(this.root);
        }
    
        /**
         *
         * @complexity `O(logn)`
         * @description Returns true if value is a member.
         */
        has(value: T): boolean {
            const compare = this.compareFn;
            const dfs = (node: TreapNode<T> | null, value: T): boolean => {
                if (node == null) return false;
                if (compare(node.value, value) === 0) return true;
                if (compare(node.value, value) < 0) return dfs(node.right, value);
                return dfs(node.left, value);
            };
    
            return dfs(this.root, value);
        }
    
        /**
         *
         * @complexity `O(logn)`
         * @description Add value to sorted set.
         */
        add(...values: T[]): this {
            const compare = this.compareFn;
            const dfs = (
                node: TreapNode<T> | null,
                value: T,
                parent: TreapNode<T>,
                direction: 'left' | 'right',
            ): void => {
                if (node == null) return;
                if (compare(node.value, value) === 0) {
                    node.count++;
                    node.pushUp();
                } else if (compare(node.value, value) > 0) {
                    if (node.left) {
                        dfs(node.left, value, node, 'left');
                    } else {
                        node.left = new TreapNode(value);
                        node.pushUp();
                    }
    
                    if (TreapNode.getFac(node.left) > node.priority) {
                        parent[direction] = node.rotateRight();
                    }
                } else if (compare(node.value, value) < 0) {
                    if (node.right) {
                        dfs(node.right, value, node, 'right');
                    } else {
                        node.right = new TreapNode(value);
                        node.pushUp();
                    }
    
                    if (TreapNode.getFac(node.right) > node.priority) {
                        parent[direction] = node.rotateLeft();
                    }
                }
                parent.pushUp();
            };
    
            values.forEach(value => dfs(this.root.left, value, this.root, 'left'));
            return this;
        }
    
        /**
         *
         * @complexity `O(logn)`
         * @description Remove value from sorted set if it is a member.
         * If value is not a member, do nothing.
         */
        delete(value: T): void {
            const compare = this.compareFn;
            const dfs = (
                node: TreapNode<T> | null,
                value: T,
                parent: TreapNode<T>,
                direction: 'left' | 'right',
            ): void => {
                if (node == null) return;
    
                if (compare(node.value, value) === 0) {
                    if (node.count > 1) {
                        node.count--;
                        node?.pushUp();
                    } else if (node.left == null && node.right == null) {
                        parent[direction] = null;
                    } else {
                        // 旋到根节点
                        if (
                            node.right == null ||
                            TreapNode.getFac(node.left) > TreapNode.getFac(node.right)
                        ) {
                            parent[direction] = node.rotateRight();
                            dfs(parent[direction]?.right ?? null, value, parent[direction]!, 'right');
                        } else {
                            parent[direction] = node.rotateLeft();
                            dfs(parent[direction]?.left ?? null, value, parent[direction]!, 'left');
                        }
                    }
                } else if (compare(node.value, value) > 0) {
                    dfs(node.left, value, node, 'left');
                } else if (compare(node.value, value) < 0) {
                    dfs(node.right, value, node, 'right');
                }
    
                parent?.pushUp();
            };
    
            dfs(this.root.left, value, this.root, 'left');
        }
    
        /**
         *
         * @complexity `O(logn)`
         * @description Returns an index to insert value in the sorted set.
         * If the value is already present, the insertion point will be before (to the left of) any existing values.
         */
        bisectLeft(value: T): number {
            const compare = this.compareFn;
            const dfs = (node: TreapNode<T> | null, value: T): number => {
                if (node == null) return 0;
    
                if (compare(node.value, value) === 0) {
                    return TreapNode.getSize(node.left);
                } else if (compare(node.value, value) > 0) {
                    return dfs(node.left, value);
                } else if (compare(node.value, value) < 0) {
                    return dfs(node.right, value) + TreapNode.getSize(node.left) + node.count;
                }
    
                return 0;
            };
    
            return dfs(this.root, value) - 1;
        }
    
        /**
         *
         * @complexity `O(logn)`
         * @description Returns an index to insert value in the sorted set.
         * If the value is already present, the insertion point will be before (to the right of) any existing values.
         */
        bisectRight(value: T): number {
            const compare = this.compareFn;
            const dfs = (node: TreapNode<T> | null, value: T): number => {
                if (node == null) return 0;
    
                if (compare(node.value, value) === 0) {
                    return TreapNode.getSize(node.left) + node.count;
                } else if (compare(node.value, value) > 0) {
                    return dfs(node.left, value);
                } else if (compare(node.value, value) < 0) {
                    return dfs(node.right, value) + TreapNode.getSize(node.left) + node.count;
                }
    
                return 0;
            };
            return dfs(this.root, value) - 1;
        }
    
        /**
         *
         * @complexity `O(logn)`
         * @description Returns the index of the first occurrence of a value in the set, or -1 if it is not present.
         */
        indexOf(value: T): number {
            const compare = this.compareFn;
            let isExist = false;
    
            const dfs = (node: TreapNode<T> | null, value: T): number => {
                if (node == null) return 0;
    
                if (compare(node.value, value) === 0) {
                    isExist = true;
                    return TreapNode.getSize(node.left);
                } else if (compare(node.value, value) > 0) {
                    return dfs(node.left, value);
                } else if (compare(node.value, value) < 0) {
                    return dfs(node.right, value) + TreapNode.getSize(node.left) + node.count;
                }
    
                return 0;
            };
            const res = dfs(this.root, value) - 1;
            return isExist ? res : -1;
        }
    
        /**
         *
         * @complexity `O(logn)`
         * @description Returns the index of the last occurrence of a value in the set, or -1 if it is not present.
         */
        lastIndexOf(value: T): number {
            const compare = this.compareFn;
            let isExist = false;
    
            const dfs = (node: TreapNode<T> | null, value: T): number => {
                if (node == null) return 0;
    
                if (compare(node.value, value) === 0) {
                    isExist = true;
                    return TreapNode.getSize(node.left) + node.count - 1;
                } else if (compare(node.value, value) > 0) {
                    return dfs(node.left, value);
                } else if (compare(node.value, value) < 0) {
                    return dfs(node.right, value) + TreapNode.getSize(node.left) + node.count;
                }
    
                return 0;
            };
    
            const res = dfs(this.root, value) - 1;
            return isExist ? res : -1;
        }
    
        /**
         *
         * @complexity `O(logn)`
         * @description Returns the item located at the specified index.
         * @param index The zero-based index of the desired code unit. A negative index will count back from the last item.
         */
        at(index: number): T | undefined {
            if (index < 0) index += this.size;
            if (index < 0 || index >= this.size) return undefined;
    
            const dfs = (node: TreapNode<T> | null, rank: number): T | undefined => {
                if (node == null) return undefined;
    
                if (TreapNode.getSize(node.left) >= rank) {
                    return dfs(node.left, rank);
                } else if (TreapNode.getSize(node.left) + node.count >= rank) {
                    return node.value;
                } else {
                    return dfs(node.right, rank - TreapNode.getSize(node.left) - node.count);
                }
            };
    
            const res = dfs(this.root, index + 2);
            return ([this.leftBound, this.rightBound] as any[]).includes(res) ? undefined : res;
        }
    
        /**
         *
         * @complexity `O(logn)`
         * @description Find and return the element less than `val`, return `undefined` if no such element found.
         */
        lower(value: T): T | undefined {
            const compare = this.compareFn;
            const dfs = (node: TreapNode<T> | null, value: T): T | undefined => {
                if (node == null) return undefined;
                if (compare(node.value, value) >= 0) return dfs(node.left, value);
    
                const tmp = dfs(node.right, value);
                if (tmp == null || compare(node.value, tmp) > 0) {
                    return node.value;
                } else {
                    return tmp;
                }
            };
    
            const res = dfs(this.root, value) as any;
            return res === this.leftBound ? undefined : res;
        }
    
        /**
         *
         * @complexity `O(logn)`
         * @description Find and return the element greater than `val`, return `undefined` if no such element found.
         */
        higher(value: T): T | undefined {
            const compare = this.compareFn;
            const dfs = (node: TreapNode<T> | null, value: T): T | undefined => {
                if (node == null) return undefined;
                if (compare(node.value, value) <= 0) return dfs(node.right, value);
    
                const tmp = dfs(node.left, value);
    
                if (tmp == null || compare(node.value, tmp) < 0) {
                    return node.value;
                } else {
                    return tmp;
                }
            };
    
            const res = dfs(this.root, value) as any;
            return res === this.rightBound ? undefined : res;
        }
    
        /**
         *
         * @complexity `O(logn)`
         * @description Find and return the element less than or equal to `val`, return `undefined` if no such element found.
         */
        floor(value: T): T | undefined {
            const compare = this.compareFn;
            const dfs = (node: TreapNode<T> | null, value: T): T | undefined => {
                if (node == null) return undefined;
                if (compare(node.value, value) === 0) return node.value;
                if (compare(node.value, value) >= 0) return dfs(node.left, value);
    
                const tmp = dfs(node.right, value);
                if (tmp == null || compare(node.value, tmp) > 0) {
                    return node.value;
                } else {
                    return tmp;
                }
            };
    
            const res = dfs(this.root, value) as any;
            return res === this.leftBound ? undefined : res;
        }
    
        /**
         *
         * @complexity `O(logn)`
         * @description Find and return the element greater than or equal to `val`, return `undefined` if no such element found.
         */
        ceil(value: T): T | undefined {
            const compare = this.compareFn;
            const dfs = (node: TreapNode<T> | null, value: T): T | undefined => {
                if (node == null) return undefined;
                if (compare(node.value, value) === 0) return node.value;
                if (compare(node.value, value) <= 0) return dfs(node.right, value);
    
                const tmp = dfs(node.left, value);
    
                if (tmp == null || compare(node.value, tmp) < 0) {
                    return node.value;
                } else {
                    return tmp;
                }
            };
    
            const res = dfs(this.root, value) as any;
            return res === this.rightBound ? undefined : res;
        }
    
        /**
         * @complexity `O(logn)`
         * @description
         * Returns the last element from set.
         * If the set is empty, undefined is returned.
         */
        first(): T | undefined {
            const iter = this.inOrder();
            iter.next();
            const res = iter.next().value;
            return res === this.rightBound ? undefined : res;
        }
    
        /**
         * @complexity `O(logn)`
         * @description
         * Returns the last element from set.
         * If the set is empty, undefined is returned .
         */
        last(): T | undefined {
            const iter = this.reverseInOrder();
            iter.next();
            const res = iter.next().value;
            return res === this.leftBound ? undefined : res;
        }
    
        /**
         * @complexity `O(logn)`
         * @description
         * Removes the first element from an set and returns it.
         * If the set is empty, undefined is returned and the set is not modified.
         */
        shift(): T | undefined {
            const first = this.first();
            if (first === undefined) return undefined;
            this.delete(first);
            return first;
        }
    
        /**
         * @complexity `O(logn)`
         * @description
         * Removes the last element from an set and returns it.
         * If the set is empty, undefined is returned and the set is not modified.
         */
        pop(index?: number): T | undefined {
            if (index == null) {
                const last = this.last();
                if (last === undefined) return undefined;
                this.delete(last);
                return last;
            }
    
            const toDelete = this.at(index);
            if (toDelete == null) return;
            this.delete(toDelete);
            return toDelete;
        }
    
        /**
         *
         * @complexity `O(logn)`
         * @description
         * Returns number of occurrences of value in the sorted set.
         */
        count(value: T): number {
            const compare = this.compareFn;
            const dfs = (node: TreapNode<T> | null, value: T): number => {
                if (node == null) return 0;
                if (compare(node.value, value) === 0) return node.count;
                if (compare(node.value, value) < 0) return dfs(node.right, value);
                return dfs(node.left, value);
            };
    
            return dfs(this.root, value);
        }
    
        *[Symbol.iterator](): Generator<T, any, any> {
            yield* this.values();
        }
    
        /**
         * @description
         * Returns an iterable of keys in the set.
         */
        *keys(): Generator<T, any, any> {
            yield* this.values();
        }
    
        /**
         * @description
         * Returns an iterable of values in the set.
         */
        *values(): Generator<T, any, any> {
            const iter = this.inOrder();
            iter.next();
            const steps = this.size;
            for (let _ = 0; _ < steps; _++) {
                yield iter.next().value;
            }
        }
    
        /**
         * @description
         * Returns a generator for reversed order traversing the set.
         */
        *rvalues(): Generator<T, any, any> {
            const iter = this.reverseInOrder();
            iter.next();
            const steps = this.size;
            for (let _ = 0; _ < steps; _++) {
                yield iter.next().value;
            }
        }
    
        /**
         * @description
         * Returns an iterable of key, value pairs for every entry in the set.
         */
        *entries(): IterableIterator<[number, T]> {
            const iter = this.inOrder();
            iter.next();
            const steps = this.size;
            for (let i = 0; i < steps; i++) {
                yield [i, iter.next().value];
            }
        }
    
        private *inOrder(root: TreapNode<T> | null = this.root): Generator<T, any, any> {
            if (root == null) return;
            yield* this.inOrder(root.left);
            const count = root.count;
            for (let _ = 0; _ < count; _++) {
                yield root.value;
            }
            yield* this.inOrder(root.right);
        }
    
        private *reverseInOrder(root: TreapNode<T> | null = this.root): Generator<T, any, any> {
            if (root == null) return;
            yield* this.reverseInOrder(root.right);
            const count = root.count;
            for (let _ = 0; _ < count; _++) {
                yield root.value;
            }
            yield* this.reverseInOrder(root.left);
        }
    }
    
    

All Problems

All Solutions