Welcome to Subscribe On Youtube
3578. Count Partitions With Max-Min Difference at Most K
Description
You are given an integer array nums
and an integer k
. Your task is to partition nums
into one or more non-empty contiguous segments such that in each segment, the difference between its maximum and minimum elements is at most k
.
Return the total number of ways to partition nums
under this condition.
Since the answer may be too large, return it modulo 109 + 7
.
Example 1:
Input: nums = [9,4,1,3,7], k = 4
Output: 6
Explanation:
There are 6 valid partitions where the difference between the maximum and minimum elements in each segment is at most k = 4
:
[[9], [4], [1], [3], [7]]
[[9], [4], [1], [3, 7]]
[[9], [4], [1, 3], [7]]
[[9], [4, 1], [3], [7]]
[[9], [4, 1], [3, 7]]
[[9], [4, 1, 3], [7]]
Example 2:
Input: nums = [3,3,4], k = 0
Output: 2
Explanation:
There are 2 valid partitions that satisfy the given conditions:
[[3], [3], [4]]
[[3, 3], [4]]
Constraints:
2 <= nums.length <= 5 * 104
1 <= nums[i] <= 109
0 <= k <= 109
Solutions
Solution 1: Dynamic Programming + Two Pointers + Ordered Set
We define $f[i]$ as the number of ways to partition the first $i$ elements. If an array satisfies that the difference between its maximum and minimum values does not exceed $k$, then any of its subarrays also satisfies this condition. Therefore, we can use two pointers to maintain a sliding window representing the current subarray.
When we reach the $r$-th element, we need to find the left pointer $l$ such that the subarray from $l$ to $r$ satisfies that the difference between the maximum and minimum values does not exceed $k$. We can use an ordered set to maintain the elements in the current window, so that we can quickly get the maximum and minimum values.
Each time we add a new element, we insert it into the ordered set and check the difference between the maximum and minimum values in the current window. If it exceeds $k$, we move the left pointer $l$ until the condition is satisfied. The number of partition ways ending at the $r$-th element is $f[l - 1] + f[l] + \ldots + f[r - 1]$. We can use a prefix sum array to quickly calculate this value.
The answer is $f[n]$, where $n$ is the length of the array.
The time complexity is $O(n \times \log n)$ and the space complexity is $O(n)$, where $n$ is the length of the array $\textit{nums}$.
-
class Solution { public int countPartitions(int[] nums, int k) { final int mod = (int) 1e9 + 7; TreeMap<Integer, Integer> sl = new TreeMap<>(); int n = nums.length; int[] f = new int[n + 1]; int[] g = new int[n + 1]; f[0] = 1; g[0] = 1; int l = 1; for (int r = 1; r <= n; r++) { int x = nums[r - 1]; sl.merge(x, 1, Integer::sum); while (sl.lastKey() - sl.firstKey() > k) { if (sl.merge(nums[l - 1], -1, Integer::sum) == 0) { sl.remove(nums[l - 1]); } ++l; } f[r] = (g[r - 1] - (l >= 2 ? g[l - 2] : 0) + mod) % mod; g[r] = (g[r - 1] + f[r]) % mod; } return f[n]; } }
-
class Solution { public: int countPartitions(vector<int>& nums, int k) { const int mod = 1e9 + 7; multiset<int> sl; int n = nums.size(); vector<int> f(n + 1, 0), g(n + 1, 0); f[0] = 1; g[0] = 1; int l = 1; for (int r = 1; r <= n; ++r) { int x = nums[r - 1]; sl.insert(x); while (*sl.rbegin() - *sl.begin() > k) { sl.erase(sl.find(nums[l - 1])); ++l; } f[r] = (g[r - 1] - (l >= 2 ? g[l - 2] : 0) + mod) % mod; g[r] = (g[r - 1] + f[r]) % mod; } return f[n]; } };
-
class Solution: def countPartitions(self, nums: List[int], k: int) -> int: mod = 10**9 + 7 sl = SortedList() n = len(nums) f = [1] + [0] * n g = [1] + [0] * n l = 1 for r, x in enumerate(nums, 1): sl.add(x) while sl[-1] - sl[0] > k: sl.remove(nums[l - 1]) l += 1 f[r] = (g[r - 1] - (g[l - 2] if l >= 2 else 0) + mod) % mod g[r] = (g[r - 1] + f[r]) % mod return f[n]
-
func countPartitions(nums []int, k int) int { const mod int = 1e9 + 7 sl := 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) } } n := len(nums) f := make([]int, n+1) g := make([]int, n+1) f[0], g[0] = 1, 1 for l, r := 1, 1; r <= n; r++ { merge(sl, nums[r-1], 1) for sl.Right().Key-sl.Left().Key > k { merge(sl, nums[l-1], -1) l++ } f[r] = g[r-1] if l >= 2 { f[r] = (f[r] - g[l-2] + mod) % mod } g[r] = (g[r-1] + f[r]) % mod } return f[n] }
-
function countPartitions(nums: number[], k: number): number { const mod = 10 ** 9 + 7; const n = nums.length; const sl = new TreapMultiSet<number>((a, b) => a - b); const f: number[] = Array(n + 1).fill(0); const g: number[] = Array(n + 1).fill(0); f[0] = 1; g[0] = 1; for (let l = 1, r = 1; r <= n; ++r) { const x = nums[r - 1]; sl.add(x); while (sl.last()! - sl.first()! > k) { sl.delete(nums[l - 1]); l++; } f[r] = (g[r - 1] - (l >= 2 ? g[l - 2] : 0) + mod) % mod; g[r] = (g[r - 1] + f[r]) % mod; } return f[n]; } 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); } }
-
use std::collections::BTreeMap; impl Solution { pub fn count_partitions(nums: Vec<i32>, k: i32) -> i32 { const mod_val: i32 = 1_000_000_007; let n = nums.len(); let mut f = vec![0; n + 1]; let mut g = vec![0; n + 1]; f[0] = 1; g[0] = 1; let mut sl = BTreeMap::new(); let mut l = 1; for r in 1..=n { let x = nums[r - 1]; *sl.entry(x).or_insert(0) += 1; while sl.keys().last().unwrap() - sl.keys().next().unwrap() > k { let val = nums[l - 1]; if let Some(cnt) = sl.get_mut(&val) { *cnt -= 1; if *cnt == 0 { sl.remove(&val); } } l += 1; } f[r] = (g[r - 1] - if l >= 2 { g[l - 2] } else { 0 } + mod_val) % mod_val; g[r] = (g[r - 1] + f[r]) % mod_val; } f[n] } }