登录[Python3/Java/C++/Go/TypeScript] 一题一解：有序集合（清晰题解）
方法一：有序集合
题目需要我们将数组 $\textit{nums}$ 分割成 $k$ 个连续且不相交的子数组，并且第二个子数组的第一个元素与第 $k$ 个子数组的第一个元素的下标距离不超过 $\textit{dist}$，这等价于让我们从 $\textit{nums}$ 下标为 $1$ 的元素开始，找到一个窗口大小为 $\textit{dist}+1$ 的子数组，求其中前 $k-1$ 的最小元素之和。我们将 $k$ 减 $1$，这样我们只需要求出 $k$ 个最小元素之和，再加上 $\textit{nums}[0]$ 即可。
我们可以使用两个有序集合 $\textit{l}$ 和 $\textit{r}$ 分别维护大小为 $\textit{dist} + 1$ 的窗口元素，其中 $\textit{l}$ 维护最小的 $k$ 个元素，而 $\textit{r}$ 维护窗口的剩余元素。我们维护一个变量 $\textit{s}$ 表示 $\textit{nums}[0]$ 与 $l$ 中元素之和。初始时，我们将前 $\textit{dist}+2$ 个元素之和累加到 $\textit{s}$ 中，并将下标为 $[1, \textit{dist} + 1]$ 的所有元素加入到 $\textit{l}$ 中。如果 $\textit{l}$ 的大小大于 $k$，我们循环将 $\textit{l}$ 中的最大元素移动到 $\textit{r}$ 中，直到 $\textit{l}$ 的大小等于 $k$，过程中更新 $\textit{s}$ 的值。
那么此时初始答案 $\textit{ans} = \textit{s}$。
接下来我们从 $\textit{dist}+2$ 开始遍历 $\textit{nums}$，对于每一个元素 $\textit{nums}[i]$，我们需要将 $\textit{nums}[i-\textit{dist}-1]$ 从 $\textit{l}$ 或 $\textit{r}$ 中移除，然后将 $\textit{nums}[i]$ 加入到 $\textit{l}$ 或 $\textit{r}$ 中。如果 $\textit{nums}[i]$ 小于 $\textit{l}$ 中的最大元素，我们将 $\textit{nums}[i]$ 加入到 $\textit{l}$ 中，否则加入到 $\textit{r}$ 中。如果此时 $\textit{l}$ 的大小小于 $k$，我们将 $\textit{r}$ 中的最小元素移动到 $\textit{l}$ 中，直到 $\textit{l}$ 的大小等于 $k$。如果此时 $\textit{l}$ 的大小大于 $k$，我们将 $\textit{l}$ 中的最大元素移动到 $\textit{r}$ 中，直到 $\textit{l}$ 的大小等于 $k$。过程中更新 $\textit{s}$ 的值，并更新 $\textit{ans} = \min(\textit{ans}, \textit{s})$。
最终答案即为 $\textit{ans}$。
###python
class Solution:
def minimumCost(self, nums: List[int], k: int, dist: int) -> int:
def l2r():
nonlocal s
x = l.pop()
s -= x
r.add(x)
def r2l():
nonlocal s
x = r.pop(0)
l.add(x)
s += x
k -= 1
s = sum(nums[: dist + 2])
l = SortedList(nums[1 : dist + 2])
r = SortedList()
while len(l) > k:
l2r()
ans = s
for i in range(dist + 2, len(nums)):
x = nums[i - dist - 1]
if x in l:
l.remove(x)
s -= x
else:
r.remove(x)
y = nums[i]
if y < l[-1]:
l.add(y)
s += y
else:
r.add(y)
while len(l) < k:
r2l()
while len(l) > k:
l2r()
ans = min(ans, s)
return ans
###java
class Solution {
private final TreeMap<Integer, Integer> l = new TreeMap<>();
private final TreeMap<Integer, Integer> r = new TreeMap<>();
private long s;
private int size;
public long minimumCost(int[] nums, int k, int dist) {
--k;
s = nums[0];
for (int i = 1; i < dist + 2; ++i) {
s += nums[i];
l.merge(nums[i], 1, Integer::sum);
}
size = dist + 1;
while (size > k) {
l2r();
}
long ans = s;
for (int i = dist + 2; i < nums.length; ++i) {
int x = nums[i - dist - 1];
if (l.containsKey(x)) {
if (l.merge(x, -1, Integer::sum) == 0) {
l.remove(x);
}
s -= x;
--size;
} else if (r.merge(x, -1, Integer::sum) == 0) {
r.remove(x);
}
int y = nums[i];
if (y < l.lastKey()) {
l.merge(y, 1, Integer::sum);
++size;
s += y;
} else {
r.merge(y, 1, Integer::sum);
}
while (size < k) {
r2l();
}
while (size > k) {
l2r();
}
ans = Math.min(ans, s);
}
return ans;
}
private void l2r() {
int x = l.lastKey();
s -= x;
if (l.merge(x, -1, Integer::sum) == 0) {
l.remove(x);
}
--size;
r.merge(x, 1, Integer::sum);
}
private void r2l() {
int x = r.firstKey();
if (r.merge(x, -1, Integer::sum) == 0) {
r.remove(x);
}
l.merge(x, 1, Integer::sum);
s += x;
++size;
}
}
###cpp
class Solution {
public:
long long minimumCost(vector<int>& nums, int k, int dist) {
--k;
multiset<int> l(nums.begin() + 1, nums.begin() + dist + 2), r;
long long s = accumulate(nums.begin(), nums.begin() + dist + 2, 0LL);
while (l.size() > k) {
int x = *l.rbegin();
l.erase(l.find(x));
s -= x;
r.insert(x);
}
long long ans = s;
for (int i = dist + 2; i < nums.size(); ++i) {
int x = nums[i - dist - 1];
auto it = l.find(x);
if (it != l.end()) {
l.erase(it);
s -= x;
} else {
r.erase(r.find(x));
}
int y = nums[i];
if (y < *l.rbegin()) {
l.insert(y);
s += y;
} else {
r.insert(y);
}
while (l.size() == k - 1) {
int x = *r.begin();
r.erase(r.find(x));
l.insert(x);
s += x;
}
while (l.size() == k + 1) {
int x = *l.rbegin();
l.erase(l.find(x));
s -= x;
r.insert(x);
}
ans = min(ans, s);
}
return ans;
}
};
###go
func minimumCost(nums []int, k int, dist int) int64 {
k--
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)
}
}
s := nums[0]
for _, x := range nums[1 : dist+2] {
s += x
merge(l, x, 1)
}
size := dist + 1
l2r := func() {
x := l.Right().Key
merge(l, x, -1)
s -= x
size--
merge(r, x, 1)
}
r2l := func() {
x := r.Left().Key
merge(r, x, -1)
merge(l, x, 1)
s += x
size++
}
for size > k {
l2r()
}
ans := s
for i := dist + 2; i < len(nums); i++ {
x := nums[i-dist-1]
if _, ok := l.Get(x); ok {
merge(l, x, -1)
s -= x
size--
} else {
merge(r, x, -1)
}
y := nums[i]
if y < l.Right().Key {
merge(l, y, 1)
s += y
size++
} else {
merge(r, y, 1)
}
for size < k {
r2l()
}
for size > k {
l2r()
}
ans = min(ans, s)
}
return int64(ans)
}
###ts
function minimumCost(nums: number[], k: number, dist: number): number {
--k;
const l = new TreapMultiSet<number>((a, b) => a - b);
const r = new TreapMultiSet<number>((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<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);
}
}
时间复杂度 $O(n \times \log \textit{dist})$，空间复杂度 $O(\textit{dist})$。其中 $n$ 为数组 $\textit{nums}$ 的长度。
有任何问题，欢迎评论区交流，欢迎评论区提供其它解题思路（代码），也可以点个赞支持一下作者哈 😄~
