diff --git a/solution/3000-3099/3013.Divide an Array Into Subarrays With Minimum Cost II/README.md b/solution/3000-3099/3013.Divide an Array Into Subarrays With Minimum Cost II/README.md index 1766ed8fda2a8..9a301a2be381f 100644 --- a/solution/3000-3099/3013.Divide an Array Into Subarrays With Minimum Cost II/README.md +++ b/solution/3000-3099/3013.Divide an Array Into Subarrays With Minimum Cost II/README.md @@ -338,6 +338,677 @@ func minimumCost(nums []int, k int, dist int) int64 { } ``` +#### TypeScript + +```ts +function minimumCost(nums: number[], k: number, dist: number): number { + --k; + const l = new TreapMultiSet((a, b) => a - b); + const r = new TreapMultiSet((a, b) => a - b); + let s = nums[0]; + for (let i = 1; i < dist + 2; ++i) { + s += nums[i]; + l.add(nums[i]); + } + const l2r = () => { + const x = l.pop()!; + s -= x; + r.add(x); + }; + const r2l = () => { + const x = r.shift()!; + l.add(x); + s += x; + }; + while (l.size > k) { + l2r(); + } + let ans = s; + for (let i = dist + 2; i < nums.length; ++i) { + const x = nums[i - dist - 1]; + if (l.has(x)) { + l.delete(x); + s -= x; + } else { + r.delete(x); + } + const y = nums[i]; + if (y < l.last()!) { + l.add(y); + s += y; + } else { + r.add(y); + } + while (l.size < k) { + r2l(); + } + while (l.size > k) { + l2r(); + } + ans = Math.min(ans, s); + } + return ans; +} + +type CompareFunction = ( + a: T, + b: T, +) => R extends 'number' ? number : boolean; + +interface ITreapMultiSet extends Iterable { + 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; + values: () => IterableIterator; + rvalues: () => IterableIterator; + entries: () => IterableIterator<[number, T]>; + + readonly size: number; +} + +class TreapNode { + value: T; + count: number; + size: number; + priority: number; + left: TreapNode | null; + right: TreapNode | 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 | null): number { + return node?.size ?? 0; + } + + static getFac(node: TreapNode | 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 { + // eslint-disable-next-line @typescript-eslint/no-this-alias + let node: TreapNode = 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 { + // eslint-disable-next-line @typescript-eslint/no-this-alias + let node: TreapNode = 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 implements ITreapMultiSet { + private readonly root: TreapNode; + private readonly compareFn: CompareFunction; + private readonly leftBound: T; + private readonly rightBound: T; + + constructor(compareFn?: CompareFunction); + constructor(compareFn: CompareFunction, leftBound: T, rightBound: T); + constructor( + compareFn: CompareFunction = (a: any, b: any) => a - b, + leftBound: any = -Infinity, + rightBound: any = Infinity, + ) { + this.root = new TreapNode(rightBound); + this.root.priority = Infinity; + this.root.left = new TreapNode(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 | 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 | 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 | null, + value: T, + parent: TreapNode, + 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 | null, + value: T, + parent: TreapNode, + 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 | 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 | 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 | 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 | 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 | 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 | 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 | 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 | 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 | 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 | 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 { + yield* this.values(); + } + + /** + * @description + * Returns an iterable of keys in the set. + */ + *keys(): Generator { + yield* this.values(); + } + + /** + * @description + * Returns an iterable of values in the set. + */ + *values(): Generator { + 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 { + 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 | null = this.root): Generator { + 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 | null = this.root): Generator { + 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); + } +} +``` + diff --git a/solution/3300-3399/3318.Find X-Sum of All K-Long Subarrays I/README.md b/solution/3300-3399/3318.Find X-Sum of All K-Long Subarrays I/README.md index 2ca2e860a26df..73c8bc89cbb34 100644 --- a/solution/3300-3399/3318.Find X-Sum of All K-Long Subarrays I/README.md +++ b/solution/3300-3399/3318.Find X-Sum of All K-Long Subarrays I/README.md @@ -76,26 +76,214 @@ edit_url: https://github.com/doocs/leetcode/edit/main/solution/3300-3399/3318.Fi -### 方法一 +### 方法一：哈希表 + 有序集合 + +我们用一个哈希表 $\textit{cnt}$ 统计窗口中每个元素的出现次数，用一个有序集合 $\textit{l}$ 存储窗口中出现次数最多的 $x$ 个元素，用另一个有序集合 $\textit{r}$ 存储剩余的元素。 + +我们维护一个变量 $\textit{s}$ 表示 $\textit{l}$ 中元素的和。初始时，我们将前 $k$ 个元素加入到窗口中，并且更新有序集合 $\textit{l}$ 和 $\textit{r}$，并且计算 $\textit{s}$ 的值。如果 $\textit{l}$ 的大小小于 $x$，并且 $\textit{r}$ 不为空，我们就循环将 $\textit{r}$ 中的最大元素移动到 $\textit{l}$ 中，直到 $\textit{l}$ 的大小等于 $x$，过程中更新 $\textit{s}$ 的值。如果 $\textit{l}$ 的大小大于 $x$，我们就循环将 $\textit{l}$ 中的最小元素移动到 $\textit{r}$ 中，直到 $\textit{l}$ 的大小等于 $x$，过程中更新 $\textit{s}$ 的值。此时，我们就可以计算出当前窗口的 $\textit{x-sum}$，添加到答案数组中。然后我们将窗口的左边界元素移出，更新 $\textit{cnt}$，并且更新有序集合 $\textit{l}$ 和 $\textit{r}$，以及 $\textit{s}$ 的值。继续遍历数组，直到遍历结束。 + +时间复杂度 $O(n \times \log k)$，空间复杂度 $O(n)$。其中 $n$ 为数组 $\textit{nums}$ 的长度。 + +相似题目： + +- [3013. 将数组分成最小总代价的子数组 II](/solution/3000-3099/3013.Divide%20an%20Array%20Into%20Subarrays%20With%20Minimum%20Cost%20II/README.md) #### Python3 ```python - +from sortedcontainers import SortedList + + +class Solution: + def findXSum(self, nums: List[int], k: int, x: int) -> List[int]: + def add(v: int): + if cnt[v] == 0: + return + p = (cnt[v], v) + if l and p > l[0]: + nonlocal s + s += p[0] * p[1] + l.add(p) + else: + r.add(p) + + def remove(v: int): + if cnt[v] == 0: + return + p = (cnt[v], v) + if p in l: + nonlocal s + s -= p[0] * p[1] + l.remove(p) + else: + r.remove(p) + + l = SortedList() + r = SortedList() + cnt = Counter() + s = 0 + n = len(nums) + ans = [0] * (n - k + 1) + for i, v in enumerate(nums): + remove(v) + cnt[v] += 1 + add(v) + j = i - k + 1 + if j < 0: + continue + while r and len(l) < x: + p = r.pop() + l.add(p) + s += p[0] * p[1] + while len(l) > x: + p = l.pop(0) + s -= p[0] * p[1] + r.add(p) + ans[j] = s + + remove(nums[j]) + cnt[nums[j]] -= 1 + add(nums[j]) + return ans ``` #### Java ```java - +class Solution { + private TreeSet l = new TreeSet<>((a, b) -> a[0] == b[0] ? a[1] - b[1] : a[0] - b[0]); + private TreeSet r = new TreeSet<>(l.comparator()); + private Map cnt = new HashMap<>(); + private int s; + + public int[] findXSum(int[] nums, int k, int x) { + int n = nums.length; + int[] ans = new int[n - k + 1]; + for (int i = 0; i < n; ++i) { + int v = nums[i]; + remove(v); + cnt.merge(v, 1, Integer::sum); + add(v); + int j = i - k + 1; + if (j < 0) { + continue; + } + while (!r.isEmpty() && l.size() < x) { + var p = r.pollLast(); + s += p[0] * p[1]; + l.add(p); + } + while (l.size() > x) { + var p = l.pollFirst(); + s -= p[0] * p[1]; + r.add(p); + } + ans[j] = s; + + remove(nums[j]); + cnt.merge(nums[j], -1, Integer::sum); + add(nums[j]); + } + return ans; + } + + private void remove(int v) { + if (!cnt.containsKey(v)) { + return; + } + var p = new int[] {cnt.get(v), v}; + if (l.contains(p)) { + l.remove(p); + s -= p[0] * p[1]; + } else { + r.remove(p); + } + } + + private void add(int v) { + if (!cnt.containsKey(v)) { + return; + } + var p = new int[] {cnt.get(v), v}; + if (!l.isEmpty() && l.comparator().compare(l.first(), p) < 0) { + l.add(p); + s += p[0] * p[1]; + } else { + r.add(p); + } + } +} ``` #### C++ ```cpp - +class Solution { +public: + vector findXSum(vector& nums, int k, int x) { + using pii = pair; + set l, r; + int s = 0; + unordered_map cnt; + auto add = [&](int v) { + if (cnt[v] == 0) { + return; + } + pii p = {cnt[v], v}; + if (!l.empty() && p > *l.begin()) { + s += p.first * p.second; + l.insert(p); + } else { + r.insert(p); + } + }; + auto remove = [&](int v) { + if (cnt[v] == 0) { + return; + } + pii p = {cnt[v], v}; + auto it = l.find(p); + if (it != l.end()) { + s -= p.first * p.second; + l.erase(it); + } else { + r.erase(p); + } + }; + vector ans; + for (int i = 0; i < nums.size(); ++i) { + remove(nums[i]); + ++cnt[nums[i]]; + add(nums[i]); + + int j = i - k + 1; + if (j < 0) { + continue; + } + + while (!r.empty() && l.size() < x) { + pii p = *r.rbegin(); + s += p.first * p.second; + r.erase(p); + l.insert(p); + } + while (l.size() > x) { + pii p = *l.begin(); + s -= p.first * p.second; + l.erase(p); + r.insert(p); + } + ans.push_back(s); + + remove(nums[j]); + --cnt[nums[j]]; + add(nums[j]); + } + return ans; + } +}; ``` #### Go diff --git a/solution/3300-3399/3318.Find X-Sum of All K-Long Subarrays I/README_EN.md b/solution/3300-3399/3318.Find X-Sum of All K-Long Subarrays I/README_EN.md index 62143c5d02775..e1423e350f648 100644 --- a/solution/3300-3399/3318.Find X-Sum of All K-Long Subarrays I/README_EN.md +++ b/solution/3300-3399/3318.Find X-Sum of All K-Long Subarrays I/README_EN.md @@ -72,26 +72,214 @@ edit_url: https://github.com/doocs/leetcode/edit/main/solution/3300-3399/3318.Fi -### Solution 1 +### Solution 1: Hash Table + Ordered Set + +We use a hash table $\textit{cnt}$ to count the occurrences of each element in the window, an ordered set $\textit{l}$ to store the $x$ elements with the highest occurrences in the window, and another ordered set $\textit{r}$ to store the remaining elements. + +We maintain a variable $\textit{s}$ to represent the sum of the elements in $\textit{l}$. Initially, we add the first $k$ elements to the window, update the ordered sets $\textit{l}$ and $\textit{r}$, and calculate the value of $\textit{s}$. If the size of $\textit{l}$ is less than $x$ and $\textit{r}$ is not empty, we repeatedly move the largest element from $\textit{r}$ to $\textit{l}$ until the size of $\textit{l}$ equals $x$, updating the value of $\textit{s}$ in the process. If the size of $\textit{l}$ is greater than $x$, we repeatedly move the smallest element from $\textit{l}$ to $\textit{r}$ until the size of $\textit{l}$ equals $x$, updating the value of $\textit{s}$ in the process. At this point, we can calculate the current window's $\textit{x-sum}$ and add it to the answer array. Then we remove the left boundary element of the window, update $\textit{cnt}$, and update the ordered sets $\textit{l}$ and $\textit{r}$, as well as the value of $\textit{s}$. Continue traversing the array until the traversal is complete. + +The time complexity is $O(n \times \log k)$, and the space complexity is $O(n)$. Here, $n$ is the length of the array $\textit{nums}$. + +Similar problems: + +- [3013. Divide an Array Into Subarrays With Minimum Cost II](/solution/3000-3099/3013.Divide%20an%20Array%20Into%20Subarrays%20With%20Minimum%20Cost%20II/README_EN.md) #### Python3 ```python - +from sortedcontainers import SortedList + + +class Solution: + def findXSum(self, nums: List[int], k: int, x: int) -> List[int]: + def add(v: int): + if cnt[v] == 0: + return + p = (cnt[v], v) + if l and p > l[0]: + nonlocal s + s += p[0] * p[1] + l.add(p) + else: + r.add(p) + + def remove(v: int): + if cnt[v] == 0: + return + p = (cnt[v], v) + if p in l: + nonlocal s + s -= p[0] * p[1] + l.remove(p) + else: + r.remove(p) + + l = SortedList() + r = SortedList() + cnt = Counter() + s = 0 + n = len(nums) + ans = [0] * (n - k + 1) + for i, v in enumerate(nums): + remove(v) + cnt[v] += 1 + add(v) + j = i - k + 1 + if j < 0: + continue + while r and len(l) < x: + p = r.pop() + l.add(p) + s += p[0] * p[1] + while len(l) > x: + p = l.pop(0) + s -= p[0] * p[1] + r.add(p) + ans[j] = s + + remove(nums[j]) + cnt[nums[j]] -= 1 + add(nums[j]) + return ans ``` #### Java ```java - +class Solution { + private TreeSet l = new TreeSet<>((a, b) -> a[0] == b[0] ? a[1] - b[1] : a[0] - b[0]); + private TreeSet r = new TreeSet<>(l.comparator()); + private Map cnt = new HashMap<>(); + private int s; + + public int[] findXSum(int[] nums, int k, int x) { + int n = nums.length; + int[] ans = new int[n - k + 1]; + for (int i = 0; i < n; ++i) { + int v = nums[i]; + remove(v); + cnt.merge(v, 1, Integer::sum); + add(v); + int j = i - k + 1; + if (j < 0) { + continue; + } + while (!r.isEmpty() && l.size() < x) { + var p = r.pollLast(); + s += p[0] * p[1]; + l.add(p); + } + while (l.size() > x) { + var p = l.pollFirst(); + s -= p[0] * p[1]; + r.add(p); + } + ans[j] = s; + + remove(nums[j]); + cnt.merge(nums[j], -1, Integer::sum); + add(nums[j]); + } + return ans; + } + + private void remove(int v) { + if (!cnt.containsKey(v)) { + return; + } + var p = new int[] {cnt.get(v), v}; + if (l.contains(p)) { + l.remove(p); + s -= p[0] * p[1]; + } else { + r.remove(p); + } + } + + private void add(int v) { + if (!cnt.containsKey(v)) { + return; + } + var p = new int[] {cnt.get(v), v}; + if (!l.isEmpty() && l.comparator().compare(l.first(), p) < 0) { + l.add(p); + s += p[0] * p[1]; + } else { + r.add(p); + } + } +} ``` #### C++ ```cpp - +class Solution { +public: + vector findXSum(vector& nums, int k, int x) { + using pii = pair; + set l, r; + int s = 0; + unordered_map cnt; + auto add = [&](int v) { + if (cnt[v] == 0) { + return; + } + pii p = {cnt[v], v}; + if (!l.empty() && p > *l.begin()) { + s += p.first * p.second; + l.insert(p); + } else { + r.insert(p); + } + }; + auto remove = [&](int v) { + if (cnt[v] == 0) { + return; + } + pii p = {cnt[v], v}; + auto it = l.find(p); + if (it != l.end()) { + s -= p.first * p.second; + l.erase(it); + } else { + r.erase(p); + } + }; + vector ans; + for (int i = 0; i < nums.size(); ++i) { + remove(nums[i]); + ++cnt[nums[i]]; + add(nums[i]); + + int j = i - k + 1; + if (j < 0) { + continue; + } + + while (!r.empty() && l.size() < x) { + pii p = *r.rbegin(); + s += p.first * p.second; + r.erase(p); + l.insert(p); + } + while (l.size() > x) { + pii p = *l.begin(); + s -= p.first * p.second; + l.erase(p); + r.insert(p); + } + ans.push_back(s); + + remove(nums[j]); + --cnt[nums[j]]; + add(nums[j]); + } + return ans; + } +}; ``` #### Go diff --git a/solution/3300-3399/3318.Find X-Sum of All K-Long Subarrays I/Solution.cpp b/solution/3300-3399/3318.Find X-Sum of All K-Long Subarrays I/Solution.cpp new file mode 100644 index 0000000000000..916039e5e080b --- /dev/null +++ b/solution/3300-3399/3318.Find X-Sum of All K-Long Subarrays I/Solution.cpp @@ -0,0 +1,64 @@ +class Solution { +public: + vector findXSum(vector& nums, int k, int x) { + using pii = pair; + set l, r; + int s = 0; + unordered_map cnt; + auto add = [&](int v) { + if (cnt[v] == 0) { + return; + } + pii p = {cnt[v], v}; + if (!l.empty() && p > *l.begin()) { + s += p.first * p.second; + l.insert(p); + } else { + r.insert(p); + } + }; + auto remove = [&](int v) { + if (cnt[v] == 0) { + return; + } + pii p = {cnt[v], v}; + auto it = l.find(p); + if (it != l.end()) { + s -= p.first * p.second; + l.erase(it); + } else { + r.erase(p); + } + }; + vector ans; + for (int i = 0; i < nums.size(); ++i) { + remove(nums[i]); + ++cnt[nums[i]]; + add(nums[i]); + + int j = i - k + 1; + if (j < 0) { + continue; + } + + while (!r.empty() && l.size() < x) { + pii p = *r.rbegin(); + s += p.first * p.second; + r.erase(p); + l.insert(p); + } + while (l.size() > x) { + pii p = *l.begin(); + s -= p.first * p.second; + l.erase(p); + r.insert(p); + } + ans.push_back(s); + + remove(nums[j]); + --cnt[nums[j]]; + add(nums[j]); + } + return ans; + } +}; diff --git a/solution/3300-3399/3318.Find X-Sum of All K-Long Subarrays I/Solution.java b/solution/3300-3399/3318.Find X-Sum of All K-Long Subarrays I/Solution.java new file mode 100644 index 0000000000000..a019114ba0b40 --- /dev/null +++ b/solution/3300-3399/3318.Find X-Sum of All K-Long Subarrays I/Solution.java @@ -0,0 +1,63 @@ +class Solution { + private TreeSet l = new TreeSet<>((a, b) -> a[0] == b[0] ? a[1] - b[1] : a[0] - b[0]); + private TreeSet r = new TreeSet<>(l.comparator()); + private Map cnt = new HashMap<>(); + private int s; + + public int[] findXSum(int[] nums, int k, int x) { + int n = nums.length; + int[] ans = new int[n - k + 1]; + for (int i = 0; i < n; ++i) { + int v = nums[i]; + remove(v); + cnt.merge(v, 1, Integer::sum); + add(v); + int j = i - k + 1; + if (j < 0) { + continue; + } + while (!r.isEmpty() && l.size() < x) { + var p = r.pollLast(); + s += p[0] * p[1]; + l.add(p); + } + while (l.size() > x) { + var p = l.pollFirst(); + s -= p[0] * p[1]; + r.add(p); + } + ans[j] = s; + + remove(nums[j]); + cnt.merge(nums[j], -1, Integer::sum); + add(nums[j]); + } + return ans; + } + + private void remove(int v) { + if (!cnt.containsKey(v)) { + return; + } + var p = new int[] {cnt.get(v), v}; + if (l.contains(p)) { + l.remove(p); + s -= p[0] * p[1]; + } else { + r.remove(p); + } + } + + private void add(int v) { + if (!cnt.containsKey(v)) { + return; + } + var p = new int[] {cnt.get(v), v}; + if (!l.isEmpty() && l.comparator().compare(l.first(), p) < 0) { + l.add(p); + s += p[0] * p[1]; + } else { + r.add(p); + } + } +} diff --git a/solution/3300-3399/3318.Find X-Sum of All K-Long Subarrays I/Solution.py b/solution/3300-3399/3318.Find X-Sum of All K-Long Subarrays I/Solution.py new file mode 100644 index 0000000000000..f885f531a96ab --- /dev/null +++ b/solution/3300-3399/3318.Find X-Sum of All K-Long Subarrays I/Solution.py @@ -0,0 +1,54 @@ +from sortedcontainers import SortedList + + +class Solution: + def findXSum(self, nums: List[int], k: int, x: int) -> List[int]: + def add(v: int): + if cnt[v] == 0: + return + p = (cnt[v], v) + if l and p > l[0]: + nonlocal s + s += p[0] * p[1] + l.add(p) + else: + r.add(p) + + def remove(v: int): + if cnt[v] == 0: + return + p = (cnt[v], v) + if p in l: + nonlocal s + s -= p[0] * p[1] + l.remove(p) + else: + r.remove(p) + + l = SortedList() + r = SortedList() + cnt = Counter() + s = 0 + n = len(nums) + ans = [0] * (n - k + 1) + for i, v in enumerate(nums): + remove(v) + cnt[v] += 1 + add(v) + j = i - k + 1 + if j < 0: + continue + while r and len(l) < x: + p = r.pop() + l.add(p) + s += p[0] * p[1] + while len(l) > x: + p = l.pop(0) + s -= p[0] * p[1] + r.add(p) + ans[j] = s + + remove(nums[j]) + cnt[nums[j]] -= 1 + add(nums[j]) + return ans diff --git a/solution/3300-3399/3321.Find X-Sum of All K-Long Subarrays II/README.md b/solution/3300-3399/3321.Find X-Sum of All K-Long Subarrays II/README.md index 1679e35c67535..6c21906ff9d08 100644 --- a/solution/3300-3399/3321.Find X-Sum of All K-Long Subarrays II/README.md +++ b/solution/3300-3399/3321.Find X-Sum of All K-Long Subarrays II/README.md @@ -78,26 +78,214 @@ edit_url: https://github.com/doocs/leetcode/edit/main/solution/3300-3399/3321.Fi -### 方法一 +### 方法一：哈希表 + 有序集合 + +我们用一个哈希表 $\textit{cnt}$ 统计窗口中每个元素的出现次数，用一个有序集合 $\textit{l}$ 存储窗口中出现次数最多的 $x$ 个元素，用另一个有序集合 $\textit{r}$ 存储剩余的元素。 + +我们维护一个变量 $\textit{s}$ 表示 $\textit{l}$ 中元素的和。初始时，我们将前 $k$ 个元素加入到窗口中，并且更新有序集合 $\textit{l}$ 和 $\textit{r}$，并且计算 $\textit{s}$ 的值。如果 $\textit{l}$ 的大小小于 $x$，并且 $\textit{r}$ 不为空，我们就循环将 $\textit{r}$ 中的最大元素移动到 $\textit{l}$ 中，直到 $\textit{l}$ 的大小等于 $x$，过程中更新 $\textit{s}$ 的值。如果 $\textit{l}$ 的大小大于 $x$，我们就循环将 $\textit{l}$ 中的最小元素移动到 $\textit{r}$ 中，直到 $\textit{l}$ 的大小等于 $x$，过程中更新 $\textit{s}$ 的值。此时，我们就可以计算出当前窗口的 $\textit{x-sum}$，添加到答案数组中。然后我们将窗口的左边界元素移出，更新 $\textit{cnt}$，并且更新有序集合 $\textit{l}$ 和 $\textit{r}$，以及 $\textit{s}$ 的值。继续遍历数组，直到遍历结束。 + +时间复杂度 $O(n \times \log k)$，空间复杂度 $O(n)$。其中 $n$ 为数组 $\textit{nums}$ 的长度。 + +相似题目： + +- [3013. 将数组分成最小总代价的子数组 II](/solution/3000-3099/3013.Divide%20an%20Array%20Into%20Subarrays%20With%20Minimum%20Cost%20II/README.md) #### Python3 ```python - +from sortedcontainers import SortedList + + +class Solution: + def findXSum(self, nums: List[int], k: int, x: int) -> List[int]: + def add(v: int): + if cnt[v] == 0: + return + p = (cnt[v], v) + if l and p > l[0]: + nonlocal s + s += p[0] * p[1] + l.add(p) + else: + r.add(p) + + def remove(v: int): + if cnt[v] == 0: + return + p = (cnt[v], v) + if p in l: + nonlocal s + s -= p[0] * p[1] + l.remove(p) + else: + r.remove(p) + + l = SortedList() + r = SortedList() + cnt = Counter() + s = 0 + n = len(nums) + ans = [0] * (n - k + 1) + for i, v in enumerate(nums): + remove(v) + cnt[v] += 1 + add(v) + j = i - k + 1 + if j < 0: + continue + while r and len(l) < x: + p = r.pop() + l.add(p) + s += p[0] * p[1] + while len(l) > x: + p = l.pop(0) + s -= p[0] * p[1] + r.add(p) + ans[j] = s + + remove(nums[j]) + cnt[nums[j]] -= 1 + add(nums[j]) + return ans ``` #### Java ```java - +class Solution { + private TreeSet l = new TreeSet<>((a, b) -> a[0] == b[0] ? a[1] - b[1] : a[0] - b[0]); + private TreeSet r = new TreeSet<>(l.comparator()); + private Map cnt = new HashMap<>(); + private long s; + + public long[] findXSum(int[] nums, int k, int x) { + int n = nums.length; + long[] ans = new long[n - k + 1]; + for (int i = 0; i < n; ++i) { + int v = nums[i]; + remove(v); + cnt.merge(v, 1, Integer::sum); + add(v); + int j = i - k + 1; + if (j < 0) { + continue; + } + while (!r.isEmpty() && l.size() < x) { + var p = r.pollLast(); + s += 1L * p[0] * p[1]; + l.add(p); + } + while (l.size() > x) { + var p = l.pollFirst(); + s -= 1L * p[0] * p[1]; + r.add(p); + } + ans[j] = s; + + remove(nums[j]); + cnt.merge(nums[j], -1, Integer::sum); + add(nums[j]); + } + return ans; + } + + private void remove(int v) { + if (!cnt.containsKey(v)) { + return; + } + var p = new int[] {cnt.get(v), v}; + if (l.contains(p)) { + l.remove(p); + s -= 1L * p[0] * p[1]; + } else { + r.remove(p); + } + } + + private void add(int v) { + if (!cnt.containsKey(v)) { + return; + } + var p = new int[] {cnt.get(v), v}; + if (!l.isEmpty() && l.comparator().compare(l.first(), p) < 0) { + l.add(p); + s += 1L * p[0] * p[1]; + } else { + r.add(p); + } + } +} ``` #### C++ ```cpp - +class Solution { +public: + vector findXSum(vector& nums, int k, int x) { + using pii = pair; + set l, r; + long long s = 0; + unordered_map cnt; + auto add = [&](int v) { + if (cnt[v] == 0) { + return; + } + pii p = {cnt[v], v}; + if (!l.empty() && p > *l.begin()) { + s += 1LL * p.first * p.second; + l.insert(p); + } else { + r.insert(p); + } + }; + auto remove = [&](int v) { + if (cnt[v] == 0) { + return; + } + pii p = {cnt[v], v}; + auto it = l.find(p); + if (it != l.end()) { + s -= 1LL * p.first * p.second; + l.erase(it); + } else { + r.erase(p); + } + }; + vector ans; + for (int i = 0; i < nums.size(); ++i) { + remove(nums[i]); + ++cnt[nums[i]]; + add(nums[i]); + + int j = i - k + 1; + if (j < 0) { + continue; + } + + while (!r.empty() && l.size() < x) { + pii p = *r.rbegin(); + s += 1LL * p.first * p.second; + r.erase(p); + l.insert(p); + } + while (l.size() > x) { + pii p = *l.begin(); + s -= 1LL * p.first * p.second; + l.erase(p); + r.insert(p); + } + ans.push_back(s); + + remove(nums[j]); + --cnt[nums[j]]; + add(nums[j]); + } + return ans; + } +}; ``` #### Go diff --git a/solution/3300-3399/3321.Find X-Sum of All K-Long Subarrays II/README_EN.md b/solution/3300-3399/3321.Find X-Sum of All K-Long Subarrays II/README_EN.md index 9701adc4cf9c9..8ddc4c29997f7 100644 --- a/solution/3300-3399/3321.Find X-Sum of All K-Long Subarrays II/README_EN.md +++ b/solution/3300-3399/3321.Find X-Sum of All K-Long Subarrays II/README_EN.md @@ -73,26 +73,214 @@ edit_url: https://github.com/doocs/leetcode/edit/main/solution/3300-3399/3321.Fi -### Solution 1 +### Solution 1: Hash Table + Ordered Set + +We use a hash table $\textit{cnt}$ to count the occurrences of each element in the window, an ordered set $\textit{l}$ to store the $x$ elements with the highest occurrences in the window, and another ordered set $\textit{r}$ to store the remaining elements. + +We maintain a variable $\textit{s}$ to represent the sum of the elements in $\textit{l}$. Initially, we add the first $k$ elements to the window, update the ordered sets $\textit{l}$ and $\textit{r}$, and calculate the value of $\textit{s}$. If the size of $\textit{l}$ is less than $x$ and $\textit{r}$ is not empty, we repeatedly move the largest element from $\textit{r}$ to $\textit{l}$ until the size of $\textit{l}$ equals $x$, updating the value of $\textit{s}$ in the process. If the size of $\textit{l}$ is greater than $x$, we repeatedly move the smallest element from $\textit{l}$ to $\textit{r}$ until the size of $\textit{l}$ equals $x$, updating the value of $\textit{s}$ in the process. At this point, we can calculate the current window's $\textit{x-sum}$ and add it to the answer array. Then we remove the left boundary element of the window, update $\textit{cnt}$, and update the ordered sets $\textit{l}$ and $\textit{r}$, as well as the value of $\textit{s}$. Continue traversing the array until the traversal is complete. + +The time complexity is $O(n \times \log k)$, and the space complexity is $O(n)$. Here, $n$ is the length of the array $\textit{nums}$. + +Similar problems: + +- [3013. Divide an Array Into Subarrays With Minimum Cost II](/solution/3000-3099/3013.Divide%20an%20Array%20Into%20Subarrays%20With%20Minimum%20Cost%20II/README_EN.md) #### Python3 ```python - +from sortedcontainers import SortedList + + +class Solution: + def findXSum(self, nums: List[int], k: int, x: int) -> List[int]: + def add(v: int): + if cnt[v] == 0: + return + p = (cnt[v], v) + if l and p > l[0]: + nonlocal s + s += p[0] * p[1] + l.add(p) + else: + r.add(p) + + def remove(v: int): + if cnt[v] == 0: + return + p = (cnt[v], v) + if p in l: + nonlocal s + s -= p[0] * p[1] + l.remove(p) + else: + r.remove(p) + + l = SortedList() + r = SortedList() + cnt = Counter() + s = 0 + n = len(nums) + ans = [0] * (n - k + 1) + for i, v in enumerate(nums): + remove(v) + cnt[v] += 1 + add(v) + j = i - k + 1 + if j < 0: + continue + while r and len(l) < x: + p = r.pop() + l.add(p) + s += p[0] * p[1] + while len(l) > x: + p = l.pop(0) + s -= p[0] * p[1] + r.add(p) + ans[j] = s + + remove(nums[j]) + cnt[nums[j]] -= 1 + add(nums[j]) + return ans ``` #### Java ```java - +class Solution { + private TreeSet l = new TreeSet<>((a, b) -> a[0] == b[0] ? a[1] - b[1] : a[0] - b[0]); + private TreeSet r = new TreeSet<>(l.comparator()); + private Map cnt = new HashMap<>(); + private long s; + + public long[] findXSum(int[] nums, int k, int x) { + int n = nums.length; + long[] ans = new long[n - k + 1]; + for (int i = 0; i < n; ++i) { + int v = nums[i]; + remove(v); + cnt.merge(v, 1, Integer::sum); + add(v); + int j = i - k + 1; + if (j < 0) { + continue; + } + while (!r.isEmpty() && l.size() < x) { + var p = r.pollLast(); + s += 1L * p[0] * p[1]; + l.add(p); + } + while (l.size() > x) { + var p = l.pollFirst(); + s -= 1L * p[0] * p[1]; + r.add(p); + } + ans[j] = s; + + remove(nums[j]); + cnt.merge(nums[j], -1, Integer::sum); + add(nums[j]); + } + return ans; + } + + private void remove(int v) { + if (!cnt.containsKey(v)) { + return; + } + var p = new int[] {cnt.get(v), v}; + if (l.contains(p)) { + l.remove(p); + s -= 1L * p[0] * p[1]; + } else { + r.remove(p); + } + } + + private void add(int v) { + if (!cnt.containsKey(v)) { + return; + } + var p = new int[] {cnt.get(v), v}; + if (!l.isEmpty() && l.comparator().compare(l.first(), p) < 0) { + l.add(p); + s += 1L * p[0] * p[1]; + } else { + r.add(p); + } + } +} ``` #### C++ ```cpp - +class Solution { +public: + vector findXSum(vector& nums, int k, int x) { + using pii = pair; + set l, r; + long long s = 0; + unordered_map cnt; + auto add = [&](int v) { + if (cnt[v] == 0) { + return; + } + pii p = {cnt[v], v}; + if (!l.empty() && p > *l.begin()) { + s += 1LL * p.first * p.second; + l.insert(p); + } else { + r.insert(p); + } + }; + auto remove = [&](int v) { + if (cnt[v] == 0) { + return; + } + pii p = {cnt[v], v}; + auto it = l.find(p); + if (it != l.end()) { + s -= 1LL * p.first * p.second; + l.erase(it); + } else { + r.erase(p); + } + }; + vector ans; + for (int i = 0; i < nums.size(); ++i) { + remove(nums[i]); + ++cnt[nums[i]]; + add(nums[i]); + + int j = i - k + 1; + if (j < 0) { + continue; + } + + while (!r.empty() && l.size() < x) { + pii p = *r.rbegin(); + s += 1LL * p.first * p.second; + r.erase(p); + l.insert(p); + } + while (l.size() > x) { + pii p = *l.begin(); + s -= 1LL * p.first * p.second; + l.erase(p); + r.insert(p); + } + ans.push_back(s); + + remove(nums[j]); + --cnt[nums[j]]; + add(nums[j]); + } + return ans; + } +}; ``` #### Go diff --git a/solution/3300-3399/3321.Find X-Sum of All K-Long Subarrays II/Solution.cpp b/solution/3300-3399/3321.Find X-Sum of All K-Long Subarrays II/Solution.cpp new file mode 100644 index 0000000000000..3082059a6a357 --- /dev/null +++ b/solution/3300-3399/3321.Find X-Sum of All K-Long Subarrays II/Solution.cpp @@ -0,0 +1,64 @@ +class Solution { +public: + vector findXSum(vector& nums, int k, int x) { + using pii = pair; + set l, r; + long long s = 0; + unordered_map cnt; + auto add = [&](int v) { + if (cnt[v] == 0) { + return; + } + pii p = {cnt[v], v}; + if (!l.empty() && p > *l.begin()) { + s += 1LL * p.first * p.second; + l.insert(p); + } else { + r.insert(p); + } + }; + auto remove = [&](int v) { + if (cnt[v] == 0) { + return; + } + pii p = {cnt[v], v}; + auto it = l.find(p); + if (it != l.end()) { + s -= 1LL * p.first * p.second; + l.erase(it); + } else { + r.erase(p); + } + }; + vector ans; + for (int i = 0; i < nums.size(); ++i) { + remove(nums[i]); + ++cnt[nums[i]]; + add(nums[i]); + + int j = i - k + 1; + if (j < 0) { + continue; + } + + while (!r.empty() && l.size() < x) { + pii p = *r.rbegin(); + s += 1LL * p.first * p.second; + r.erase(p); + l.insert(p); + } + while (l.size() > x) { + pii p = *l.begin(); + s -= 1LL * p.first * p.second; + l.erase(p); + r.insert(p); + } + ans.push_back(s); + + remove(nums[j]); + --cnt[nums[j]]; + add(nums[j]); + } + return ans; + } +}; diff --git a/solution/3300-3399/3321.Find X-Sum of All K-Long Subarrays II/Solution.java b/solution/3300-3399/3321.Find X-Sum of All K-Long Subarrays II/Solution.java new file mode 100644 index 0000000000000..7047df89dbe5b --- /dev/null +++ b/solution/3300-3399/3321.Find X-Sum of All K-Long Subarrays II/Solution.java @@ -0,0 +1,63 @@ +class Solution { + private TreeSet l = new TreeSet<>((a, b) -> a[0] == b[0] ? a[1] - b[1] : a[0] - b[0]); + private TreeSet r = new TreeSet<>(l.comparator()); + private Map cnt = new HashMap<>(); + private long s; + + public long[] findXSum(int[] nums, int k, int x) { + int n = nums.length; + long[] ans = new long[n - k + 1]; + for (int i = 0; i < n; ++i) { + int v = nums[i]; + remove(v); + cnt.merge(v, 1, Integer::sum); + add(v); + int j = i - k + 1; + if (j < 0) { + continue; + } + while (!r.isEmpty() && l.size() < x) { + var p = r.pollLast(); + s += 1L * p[0] * p[1]; + l.add(p); + } + while (l.size() > x) { + var p = l.pollFirst(); + s -= 1L * p[0] * p[1]; + r.add(p); + } + ans[j] = s; + + remove(nums[j]); + cnt.merge(nums[j], -1, Integer::sum); + add(nums[j]); + } + return ans; + } + + private void remove(int v) { + if (!cnt.containsKey(v)) { + return; + } + var p = new int[] {cnt.get(v), v}; + if (l.contains(p)) { + l.remove(p); + s -= 1L * p[0] * p[1]; + } else { + r.remove(p); + } + } + + private void add(int v) { + if (!cnt.containsKey(v)) { + return; + } + var p = new int[] {cnt.get(v), v}; + if (!l.isEmpty() && l.comparator().compare(l.first(), p) < 0) { + l.add(p); + s += 1L * p[0] * p[1]; + } else { + r.add(p); + } + } +} diff --git a/solution/3300-3399/3321.Find X-Sum of All K-Long Subarrays II/Solution.py b/solution/3300-3399/3321.Find X-Sum of All K-Long Subarrays II/Solution.py new file mode 100644 index 0000000000000..f885f531a96ab --- /dev/null +++ b/solution/3300-3399/3321.Find X-Sum of All K-Long Subarrays II/Solution.py @@ -0,0 +1,54 @@ +from sortedcontainers import SortedList + + +class Solution: + def findXSum(self, nums: List[int], k: int, x: int) -> List[int]: + def add(v: int): + if cnt[v] == 0: + return + p = (cnt[v], v) + if l and p > l[0]: + nonlocal s + s += p[0] * p[1] + l.add(p) + else: + r.add(p) + + def remove(v: int): + if cnt[v] == 0: + return + p = (cnt[v], v) + if p in l: + nonlocal s + s -= p[0] * p[1] + l.remove(p) + else: + r.remove(p) + + l = SortedList() + r = SortedList() + cnt = Counter() + s = 0 + n = len(nums) + ans = [0] * (n - k + 1) + for i, v in enumerate(nums): + remove(v) + cnt[v] += 1 + add(v) + j = i - k + 1 + if j < 0: + continue + while r and len(l) < x: + p = r.pop() + l.add(p) + s += p[0] * p[1] + while len(l) > x: + p = l.pop(0) + s -= p[0] * p[1] + r.add(p) + ans[j] = s + + remove(nums[j]) + cnt[nums[j]] -= 1 + add(nums[j]) + return ans