Problem: Maximizing the Sum of Squared Segment Totals
You are given a list of integers and an integer k. You must split the list into exactly k non-empty consecutive parts. The “score” of a partition is calculated by taking each part, summing its numbers, then squaring that sum, and finally adding up all these squared values. Your goal is to find the maximum possible score over all ways to split the list into k consecutive parts.
Input
- On the first line: two integers, n (the length of the list) and k (the number of parts), with 1 ≤ k ≤ n ≤ 200 000.
- On the second line: n integers (each between –10000 and 10000), representing the list.
Output
- A single integer: the highest score you can achieve by splitting the list into exactly k consecutive non-empty parts.
Example
Input
5 2
1 2 -1 2 3
Output
29
How the example works
One optimal way is:
- First part takes the first three numbers: 1, 2, –1. Their sum is 2; its square is 4.
- Second part takes the last two numbers: 2, 3. Their sum is 5; its square is 25.
- Total score is 4 + 25 = 29.
Constraints
- List length n up to 200 000, so an O(n² × k) approach is too slow.
- You must handle negative and positive numbers.
- Time limit: roughly 2 seconds.
- Memory limit: 512 MB.
High-Level Solution Approach
Prefix sums for fast range sums
Build an array of running totals so that you can get the sum of any consecutive block in constant time.Dynamic programming over splits
Let dp[j][i] be the maximum score you can get by splitting the first i elements into j parts. Your final answer is dp[k][n].Transition in words
To compute dp[j][i], consider where the last split happens. Suppose the previous split ended at position t (which ranges from j–1 up to i–1). Then you know dp[j–1][t], and you know the sum of the block from t+1 to i. Squaring that block’s sum and adding it to dp[j–1][t] gives one candidate total. You take the maximum over all possible t.Speeding up with line-container trick
The brute-force inner maximization would be quadratic, but by rewriting the squared-sum expression you can treat each possible split point as a “line” (with a slope and intercept) and each query (for a new end position) as asking for the maximum value at a given x-coordinate. A data structure known as a “Li Chao tree” or a convex-hull-trick container lets you insert lines and query maximum values in logarithmic time, reducing each DP layer to roughly O(n log range_of_sums). Overall complexity becomes O(k × n log n), which is acceptable for the given limits.
JavaScript Implementation
/**
* Compute the maximum possible score by splitting the array into k parts.
* @param {number[]} arr - zero-based array of length n
* @param {number} k - number of parts to split into
* @returns {bigint} - maximum achievable score
*/
function maxPartitionScore(arr, k) {
const n = arr.length;
// Build prefix sums (using BigInt for safety)
const prefix = Array(n + 1).fill(0n);
for (let i = 1; i <= n; i++) {
prefix[i] = prefix[i - 1] + BigInt(arr[i - 1]);
}
// dpPrev and dpCurr store the DP values for the previous and current layer
let dpPrev = Array(n + 1).fill(-1n << 60n);
let dpCurr = Array(n + 1).fill(-1n << 60n);
dpPrev[0] = 0n;
// Li Chao tree implementation for maximum of lines
class LiChao {
constructor(minX, maxX) {
this.minX = minX;
this.maxX = maxX;
this.line = null;
this.left = null;
this.right = null;
}
eval(line, x) {
return line.m * x + line.b;
}
insert(newLine, left = this.minX, right = this.maxX) {
if (!this.line) {
this.line = newLine;
return;
}
const mid = (left + right) >> 1n;
// Keep whichever line is better at mid
if (this.eval(newLine, mid) > this.eval(this.line, mid)) {
[this.line, newLine] = [newLine, this.line];
}
if (right - left === 0n) return;
if (this.eval(newLine, left) > this.eval(this.line, left)) {
if (!this.left) this.left = new LiChao(left, mid);
this.left.insert(newLine, left, mid);
} else if (this.eval(newLine, right) > this.eval(this.line, right)) {
if (!this.right) this.right = new LiChao(mid + 1n, right);
this.right.insert(newLine, mid + 1n, right);
}
}
query(x, left = this.minX, right = this.maxX) {
if (!this.line) return -1n << 60n;
let result = this.eval(this.line, x);
if (left === right) return result;
const mid = (left + right) >> 1n;
if (x <= mid && this.left) {
const cand = this.left.query(x, left, mid);
if (cand > result) result = cand;
}
if (x > mid && this.right) {
const cand = this.right.query(x, mid + 1n, right);
if (cand > result) result = cand;
}
return result;
}
}
const MIN_X = -10n ** 14n;
const MAX_X = 10n ** 14n;
// Main DP loop
for (let part = 1; part <= k; part++) {
const cht = new LiChao(MIN_X, MAX_X);
// Insert the line corresponding to splitting at position 0
cht.insert({
m: -2n * prefix[0],
b: dpPrev[0] + prefix[0] * prefix[0]
});
for (let i = 1; i <= n; i++) {
// Query the best previous split for ending at i
const best = cht.query(prefix[i]);
// Score is square of prefix[i] plus best
dpCurr[i] = prefix[i] * prefix[i] + best;
// Insert the line for future queries
cht.insert({
m: -2n * prefix[i],
b: dpPrev[i] + prefix[i] * prefix[i]
});
}
// Rotate arrays for next layer
[dpPrev, dpCurr] = [dpCurr, dpPrev];
dpCurr.fill(-1n << 60n);
}
return dpPrev[n];
}
// Example
const input = [1, 2, -1, 2, 3];
const k = 2;
console.log(String(maxPartitionScore(input, k))); // prints 29
Notes on the implementation
- We use JavaScript’s
BigInttype to avoid overflow when squaring large sums. - Each layer of the DP inserts n lines into the convex-hull structure and performs n queries, giving roughly O(n log range_of_sums) per layer.
- The overall time complexity is approximately O(k × n log n), which works for large n (up to 200 000) and moderate k.
Top comments (2)
pretty sick seeing optimization tricks actually get put to work for stuff like this - you ever feel like pulling off performance like that is more about instinct or just pure grinding through math?
Great I want more