题解：P17245 [IOI 2026] 划分 / Partition

猜结论题，感觉根本不是给人做的。

设 Ulug'bek 想要每一组的元素和都为 \(V\)。一些想法大概是尝试把 \(a\) 划分成 \(K\) 组，使得每个组内的元素和与 \(V\) 之间的差距比较小。考虑一些变态的想法，我们当然希望 \(V-sum\leq M\)，我们加强一下，使得 \(V-sum\leq a_i\leq M\)。这个形式启发我们做经典的贪心：把 \(a\) 从小到大排序，每次尝试把 \(a_i\) 加入当前组中，若装不下就新开一组。观察到这样贪心后，对于某个不在末尾的组 \(a[l..r]\)，我们有 \(V-sum<a_{r+1}\leq M\)。

设 \(sum_i\) 表示第 \(i\) 组的元素和，若不足 \(K\) 组则补充 \(sum=0\) 的空组。那么最少要添加 \(\sum_{i=1}^K\left\lceil\dfrac{V-sum_i}{M}\right\rceil\) 个数。我们要尽量最小化添加的数的个数，于是考虑取最小的 \(V\) 使得分出来的组数 \(\leq K\)。我们证明，取最小的 \(V\) 时，添加的数的个数不超过 \(K-1\)。

证明

反证，假设 \(\sum_{i=1}^K\left\lceil\dfrac{V-sum_i}{M}\right\rceil\geq K\)。

考察令 \(V\gets V-1\) 时会发生什么。设到达原来的第 \(i\) 组时有 \(p_i\) 个元素因为 \(V\) 减小而被挤到当前组。

设 \(c_i=\left\lceil\dfrac{V-sum_i}{M}\right\rceil\)，则 \((c_i-1)M\leq V-1-sum_i\)。若 \(p_i\leq c_i-1\)，则被挤过来的元素都能塞进第 \(i\) 组里，于是 \(p_{i+1}=0\)；而若 \(p_i\geq c_i\)，设第 \(i\) 组原本有 \(s_i\) 个元素，那么此时第 \(i\) 组至少能容纳 \(s_i+c_i-1\) 个元素，因此 \(p_{i+1}\leq p_i+1-c_i\)。综上可得 \(p_{i+1}\leq \max(p_i+1-c_i,0)\)。

利用贪心的性质，设容量为 \(V\) 时分出了 \(x\) 组，那么对于 \(1\leq i<x\)，有 \(1-c_i\geq 0\)，对于 \(x<i\leq K\)，有 \(1-c_i\leq 0\)，因此实际上上面的限制可以写作

\[p_{i+1}\leq\max\left(i-\sum_{j=1}^ic_j,0\right) \]

由最开始的假设，可以得到 \(p_{K+1}=0\)。因此令 \(V\gets V-1\) 后仍然可以分成至多 \(K\) 组，与 \(V\) 的最小性矛盾。

注意到合法的 \(V\) 满足 \(V\geq \max a_i\)，所以这里并不完备，还需要讨论 \(V=\max a_i\) 的情况。由于 \(V\leq M\)，此时必然有 \(c_i\leq 1\)。而 \(a\) 中任意一个最大值所在的组必然有 \(c_i=0\)，因此 \(\sum c_i\leq K-1\)。\(\Box\)

直接在 \([\max a_i,\sum a_i]\) 范围内二分即可得到最小的 \(V\)。至此我们得到了一个性质很好的 Ulug'bek 的策略。

接下来考虑 Temur 如何分组。

回顾 Ulug'bek 的分组有什么性质：设填数前分出来的组为 \(S_1,\cdots,S_K\)，对于组 \(i\) 满足 \(S_{i+1}\) 非空，\(V-sum_i<\min S_{i+1}\)。

一个自然的想法是倒着做，把 \(b\) 从大到小排序，然后做相同的贪心策略。但是手玩一下会发现这个错完了，一个组内的元素完全可以不在一段连续的区间内。考虑微调一下：预先创建 \(K\) 个空组，然后每个元素放到编号最小的能放进的组里面。

感觉上这个策略挺对的，我们来证明一下。

证明

设 \(S_i\) 填数后变为 \(T_i\)。我们声称上述策略恢复出来的分组方案恰好是 \(T_K,\cdots,T_1\)。

对于原本为空的组 \(S_i\)，设 \(V=qM+r\)，那么 \(b\) 最开头的若干个 \(M\) 恰好会被按 \(i\) 倒序装进 \(T_i\) 中，接下来 \(T_i\) 将一直装不进数，直到扫到第一个 \(r\)。

对于不为空的组 \(S_i\)，考察 \(j<i\)，显然 \(S_j\) 中的元素都不超过 \(\min S_i\)，同时，\(V-sum_j<\min S_{j+1}\leq\min S_i\)。因此 \(\geq\min S_i\) 的数一定全部属于 \(T_i\)，此后 \(T_i\) 同样一直装不进数，直到扫到第一个 \(V-sum_i\)。\(\Box\)

容易线段树二分维护这个贪心的过程。于是我们也成功得到了 Temur 的策略。

第一部分时间复杂度为 \(\mathcal{O}(N\log{N}+N\log{M})\)，第二部分时间复杂度为 \(\mathcal{O}((N+K)\log{K})\)。

代码显然很好写。

主要代码

vector<int> add_numbers(vector<int> A, int K, int M) {
	auto calc = [&](ll mid) {
		int res = 1;
		ll sum = 0;
		for (int x : A) {
			sum += x;
			if (sum > mid) {
				sum = x;
				++res;
			}
		}
		return res;
	};

	sort(A.begin(), A.end());
	ll l = A.back(), r = accumulate(A.begin(), A.end(), 0ll);
	while (l < r) {
		ll mid = l + r >> 1;
		if (calc(mid) <= K) r = mid;
		else l = mid + 1;
	}

	ll V = l, sum = 0;
	vector<ll> s;
	s.reserve(K);
	for (int x : A) {
		if (sum + x > V) {
			s.emplace_back(sum);
			sum = 0;
		}
		sum += x;
	}
	s.emplace_back(sum);
	while (s.size() < K) s.emplace_back(0);

	vector<int> res;
	for (ll x : s) {
		while (x + M <= V) {
			res.emplace_back(M);
			x += M;
		}
		if (x != V) res.emplace_back(V - x);
	}
	return res;
}

struct SegTree {
#define ls(p) (p << 1)
#define rs(p) (p << 1 | 1)
	ll mx[MAXN << 2];

	void pushUp(int p) {
		mx[p] = max(mx[ls(p)], mx[rs(p)]);
	}

	void build(int p, int l, int r, ll v) {
		mx[p] = v;
		if (l == r) return;
		int mid = l + r >> 1;
		build(ls(p), l, mid, v);
		build(rs(p), mid + 1, r, v);
	}

	int find(int p, int l, int r, ll v) {
		if (mx[p] < v) return inf;
		if (l == r) {
			mx[p] -= v;
			return l;
		}
		int mid = l + r >> 1;
		int res = find(ls(p), l, mid, v);
		if (res == inf) res = find(rs(p), mid + 1, r, v);
		pushUp(p);
		return res;
	}
#undef ls
#undef rs
} sgt;

vector<int> find_partition(vector<int> B, int K) {
	ll sum = accumulate(B.begin(), B.end(), 0ll);
	sgt.build(1, 0, K - 1, sum / K);
	int N = B.size();
	vector<int> res(N);
	for (int i = N - 1; i >= 0; --i) res[i] = sgt.find(1, 0, K - 1, B[i]);
	return res;
}