| comments | difficulty | edit_url | tags | ||||
|---|---|---|---|---|---|---|---|
true |
Hard |
|
The 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 values.
- For examples, if
arr = [2,3,4], the median is3. - For examples, if
arr = [1,2,3,4], the median is(2 + 3) / 2 = 2.5.
You are given an integer array nums and an integer k. 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.
Return the median array for each window in the original array. Answers within 10-5 of the actual value will be accepted.
Example 1:
Input: nums = [1,3,-1,-3,5,3,6,7], k = 3 Output: [1.00000,-1.00000,-1.00000,3.00000,5.00000,6.00000] Explanation: 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
Example 2:
Input: nums = [1,2,3,4,2,3,1,4,2], k = 3 Output: [2.00000,3.00000,3.00000,3.00000,2.00000,3.00000,2.00000]
Constraints:
1 <= k <= nums.length <= 105-231 <= nums[i] <= 231 - 1
We can use two priority queues (min-heap and max-heap) to maintain the elements in the current window. One priority queue stores the smaller half of the elements, and the other priority queue stores the larger half of the elements. This way, the median of the current window is either the average of the top elements of the two heaps or one of the top elements.
We design a class
-
add_num(num): Adds$\textit{num}$ to the current window. -
find_median(): Returns the median of the elements in the current window. -
remove_num(num): Removes$\textit{num}$ from the current window. -
prune(pq): If the top element of the heap is in the lazy deletion dictionary$\textit{delayed}$ , it pops the top element from the heap and decrements its lazy deletion count. If the lazy deletion count of the element becomes zero, it removes the element from the lazy deletion dictionary. -
rebalance(): If the number of elements in the smaller half exceeds the number of elements in the larger half by$2$ , it moves the top element of the larger half to the smaller half. If the number of elements in the smaller half is less than the number of elements in the larger half, it moves the top element of the larger half to the smaller half.
In the add_num(num) method, we first consider adding rebalance() method to ensure that the size difference between the two priority queues does not exceed
In the remove_num(num) method, we increment the lazy deletion count of prune() method to ensure that the top element of the smaller half is not in the lazy deletion dictionary. Otherwise, we update the size of the larger half and call the prune() method to ensure that the top element of the larger half is not in the lazy deletion dictionary.
In the find_median() method, if the current window size is odd, we return the top element of the smaller half; otherwise, we return the average of the top elements of the smaller half and the larger half.
In the prune(pq) method, if the top element of the heap is in the lazy deletion dictionary, it pops the top element from the heap and decrements its lazy deletion count. If the lazy deletion count of the element becomes zero, it removes the element from the lazy deletion dictionary.
In the rebalance() method, if the number of elements in the smaller half exceeds the number of elements in the larger half by
The time complexity is
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 ansclass 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 MedianFinder {
public:
MedianFinder(int k) {
this->k = k;
}
void addNum(int num) {
if (small.empty() || num <= small.top()) {
small.push(num);
++smallSize;
} else {
large.push(num);
++largeSize;
}
reblance();
}
void removeNum(int num) {
++delayed[num];
if (num <= small.top()) {
--smallSize;
if (num == small.top()) {
prune(small);
}
} else {
--largeSize;
if (num == large.top()) {
prune(large);
}
}
reblance();
}
double findMedian() {
return k & 1 ? small.top() : ((double) small.top() + large.top()) / 2.0;
}
private:
priority_queue<int> small;
priority_queue<int, vector<int>, greater<int>> large;
unordered_map<int, int> delayed;
int smallSize = 0;
int largeSize = 0;
int k;
template <typename T>
void prune(T& pq) {
while (!pq.empty() && delayed[pq.top()]) {
if (--delayed[pq.top()] == 0) {
delayed.erase(pq.top());
}
pq.pop();
}
}
void reblance() {
if (smallSize > largeSize + 1) {
large.push(small.top());
small.pop();
--smallSize;
++largeSize;
prune(small);
} else if (smallSize < largeSize) {
small.push(large.top());
large.pop();
++smallSize;
--largeSize;
prune(large);
}
}
};
class Solution {
public:
vector<double> medianSlidingWindow(vector<int>& nums, int k) {
MedianFinder finder(k);
for (int i = 0; i < k; ++i) {
finder.addNum(nums[i]);
}
vector<double> ans = {finder.findMedian()};
for (int i = k; i < nums.size(); ++i) {
finder.addNum(nums[i]);
finder.removeNum(nums[i - k]);
ans.push_back(finder.findMedian());
}
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 any) { h.IntSlice = append(h.IntSlice, v.(int)) }
func (h *hp) Pop() any {
a := h.IntSlice
v := a[len(a)-1]
h.IntSlice = a[:len(a)-1]
return v
}We can use two ordered sets to maintain the elements in the current window. The ordered set
We traverse the array
If the total number of elements in the current window is
The time complexity is
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 ansclass 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;
}
};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);
}
}