P8226 「Wdoi-5」樱点收集 题解

要是赛时对 pb_ds 有掌握就 $200$ 分了！

提供一个平衡树写法，虽然平衡树我并不会写，但是 pbds 内置有平衡树啊，所以我们考虑用内置的平衡树写，注意要用 rb_tree_tag 而非 splay_tree_tag，内置的 Splay 很慢，开了 O2 还是超时，而内置红黑树不开 O2 都能过。

说了用平衡树，但是具体什么意思呢？我们考虑记录前缀和，对于每一个 $sum_i$ 我们记录 $sum_i$ 模 $k$ 的余数出现在哪些位置，那么显然对于每一个地方 $i$ 放炸弹他的答案其实就是： $\sum_{j=1}^{i-1} \begin{cases} 1 & sum_j \equiv 0 \pmod k \\ 0 & \text{otherwise}\end{cases} + \sum_{j=i+1}^n \begin{cases} 1 & sum_j \equiv a_i \pmod k \\ 0 & \text{otherwise} \end{cases}$。

显然这两个东西可以平衡树计算，总复杂度 $O(n \log n)$，事实上你会发现这个东西可以值域线段树，常数更小，不过这样也能过。

代码：

#include <bits/stdc++.h>
#include <bits/extc++.h>
using namespace std;
using namespace __gnu_pbds; 

#define ll long long

constexpr int N(3e5 + 5);

tree<ll, null_type, less<ll>, rb_tree_tag, tree_order_statistics_node_update> tr[N];

int n, m;
ll k, ans = 0;

ll a[N], sum[N], b[N];

int main()
{
	scanf("%d%d%lld", &n, &m, &k);
	for (register int i(1); i <= m; ++i)
	{
		scanf("%lld", &b[i]);
	}
	for (register int i(1); i <= n; ++i)
	{
		scanf("%lld", &a[i]);
		sum[i] = sum[i - 1] + a[i];
	}
	for (register int i(1); i <= m; ++i)
	{
		tr[sum[b[i]] % k].insert(b[i]);
	}
	for (register ll i(1); i <= n; ++i)
	{
		ans = max(ans, (ll)(tr[0].order_of_key(i) - tr[0].order_of_key(1) + (tr[a[i] % k].order_of_key(n + 1) - tr[a[i] % k].order_of_key(i + 1))));
	}
	ans = max(ans, (ll)(tr[0].order_of_key(n + 1) - tr[0].order_of_key(1)));
	printf("%lld\n", ans);
	return 0;
}