CF1659 D Reverse Sort Sum 树状数组

分析

设数组的长度为 \(n\)，数组 \(C\) 中所有元素之和为 \(sum\)，每个数组中 \(1\) 的数量为 \(k\)。易知 \(k = \dfrac {sum} n\)。

由定义，数组 \(B_i\) 的前 \(i\) 项是严格非递减的。再考虑到 \(c_i = \sum_{j = 1}^n B_{j, i}\)。可以发现当 \(i > 1\) 时

\[a_i = \begin{cases} 1, & \sum_{j = 1}^{i - 1} B_{j, i} = i - 1\\ 0, & \sum_{j = 1}^{i - 1} B_{j, i} = 0 \end{cases} \]

至于 \(a_1\) 值，当其余值确定时，也可唯一确定。

由上式可以发现，针对每个 \(a_i\)，我们只需要知道数组 \(B_1, B_2, \dots,B_{i - 1}\) 的前 \(i\) 项的值即可。

因为已知的是 \(B_1 \sim B_n\) 每一项的和，可以逆序依次求出 \(a_n, a_{n - 1}, \dots\)。

首先求 \(a_n\)，只要知道 \(B_1, B_2, \dots,B_{n - 1}\) 中的第 \(n\) 项的值，而可以确定的 \(B_n\) 显然在后续的计算中会产生影响，需要减去 \(C\) 中相应的值。定义此时 \(a\) 中还未确定的数中 \(1\) 的个数为 \(x\)，也就是说 \(B_n\) 的数字 \(1\) 所在的区间为 \([n + 1 - x, n - 1]\)。所以只需要对 \(C\) 中的 \([n + 1 - x, n]\) 减一即可。此时的 \(c_n\) 就等于 \(B_1, B_2, \dots,B_{n - 1}\) 中的第 \(n\) 项的值。\(a_n\) 可以确定。

类似的，接下来逐个求 \(a_i\) 的值也是类似的步骤。令此时 \(a\) 中还未确定的数中 \(1\) 的个数为 \(x\)，也就是说在区间 \([1, i]\) 中 \(1\) 的个数为 \(x\)。可知 \(B_{i}\) 中的前 \(i\) 项的后 \(x\) 项均为 \(1\)。由于\(B_{i}\) 中 \([i + 1, n]\) 对前 \(i\) 项的求解无关，无须考虑，只有前 \(i\) 项有影响，且其中 \(1\) 所在的区间为 \([i + 1 - x, i]\)，只需要对 \(C\) 中的 \([i + 1 - x, i]\) 减一即可。可以知道此时 \(c_i = \sum_{j = 1}^{i - 1} B_{j, i}\)，\(a_i\) 也可以被确定了。

其中对数组 \(C\) 实现区间修改，单点查询用树状数组维护一下就行。复杂度 \(O(N\log N)\)。

代码

#include <bits/stdc++.h>
#define rep(i, a, b) for (ll i = ll(a); i <= ll(b); ++i)

using namespace std;
using ll = long long;

const int N = 2e5 + 10;
int n, a[N], c[N];
int tr[N];

inline int lowbit(int x) {
	return x & (-x);
}

void update(int x, int d) {
	while (x <= n) {
		tr[x] += d;
		x += lowbit(x);
	}
}

int query(int x) {
	int res = 0;
	while (x > 0) {
		res += tr[x];
		x -= lowbit(x);
	}return res;
}

void solve() {
	cin >> n;
	ll sum = 0;							// 注意数组c中所有元素之和超过了int的范围
	rep (i, 1, n) {
		cin >> c[i]; sum += c[i];
	}
	int k = sum / n; 
	int x = k;
	memset(tr, 0, sizeof tr);
	for (int i = 1; i <= n; i++) {
		update(i, c[i] - c[i - 1]); 	// 对数组c的差分数组建树状数组，用来进行区间修改
	}
	for (int i = n; i >= 2; i--) {  	// 从n到2依次确定a_i的值
		if (x > 0) {
			int l = i + 1 - x, r = i;	// 将区间[i + 1 - x, i] 减一
			update(l, -1); 		
			update(r + 1, 1);
		}
		int num = query(i);				// 查询当前 c_i的值，如果为0，则a_i = 0，否则为1
		if (num == 0) a[i] = 0;
		else a[i] = 1, x--;	
	}
	a[1] = x;							// 最后剩余的x即为a_1
	rep (i, 1, n) cout << a[i] << ' ';
	cout << '\n';
}

int main()
{
    ios::sync_with_stdio(0);
    cin.tie(0);
	int _; cin >> _;
    while (_--) {
        solve();
    }
    return 0;
}