DSU on Tree 求条

#include <bits/stdc++.h>
using namespace std;
const int N = 2e5 + 10;
typedef long long ll;
typedef pair<int, int> pi;
typedef unsigned long long ull;
#define endl '\n'
int n, p, m;
struct node {
	ll hs[N], fp[N], pri, mod;
	void init (int a[], int x, int y) {
		fp[0] = 1;
		pri = x, mod = y;
		for (int i = 1; i <= n; i ++) {
			hs[i] = (hs[i - 1] * pri % mod + a[i]) % mod;
			fp[i] = fp[i - 1] * pri % mod;
		}
	}
	ll f (int l, int r) {
		return (hs[r] - hs[l - 1] * fp[r - l + 1] % mod + mod) % mod;
	}
}h, h1, h2;
int a[N];
int cnt[50000000];

signed main() {
	freopen("collision.in", "r", stdin);
	freopen("collision.out", "w", stdout);
	ios::sync_with_stdio(0), cin.tie(0), cout.tie(0);
	cin >> n >> p >> m;
	for (int i = 1; i <= n; i ++) cin >> a[n - i + 1];
	h.init(a, p, m);
	h1.init(a, 131, 998244853);
	h2.init(a, 131, 998244353);
	if (m <= 5e7) {
		ll ans = 0;
		for (int len = 1; len <= n; len ++) {
			for (int i = 1, j = len; j <= n; i ++, j ++) {
				ans += cnt[h.f(i, j)], cnt[h.f(i, j)] ++;
			}
			vector<pi> v;
			for (int i = 1, j = len; j <= n; i ++, j ++)
				v.push_back({h1.f(i, j), h2.f(i, j)});
			sort (v.begin(), v.end());
			pi lst = {-1, -1};
			int r = 0;
			for (pi x : v) {
				if (x != lst) ans -= (r - 1) * r / 2, lst = x, r = 1;
				else r ++;
			}
			ans -= (r - 1) * r / 2;
		}
		cout << ans * 2 << endl;
	} else {
		ll ans = 0;
		vector<int> ve;
		for (int len = 1; len <= n; len ++) {
			for (int i = 1, j = len; j <= n; i ++, j ++) ve.push_back(h.f(i, j));
			vector<pi> v;
			for (int i = 1, j = len; j <= n; i ++, j ++)
				v.push_back({h1.f(i, j), h2.f(i, j)});
			sort (v.begin(), v.end());
			pi lst = {-1, -1};
			int r = 0;
			for (pi x : v) {
				if (x != lst) ans -= (r - 1) * r / 2, lst = x, r = 1;
				else r ++;
			}
			ans -= (r - 1) * r / 2;
		}
		sort (ve.begin(), ve.end());
		int lst = -1, r = 0;
		for (int x : ve) {
			if (x != lst) ans += 1ll * (r - 1) * r / 2, lst = x, r = 1;
			else r ++;
		}
		ans += 1ll * (r - 1) * r / 2;
		cout << ans * 2 << endl;
	}
	return 0;
}