#include <iostream>
#include <cmath>
#include <vector>
#include <algorithm>
#include <map>
#include <iomanip>
#include <stack>
#include <set>
#include <queue>
#include <cstring>
#define endl '\n'
#define int long long
template<typename T1, typename T2> T1 chkmin (T1 x, T2 y) {return x < y ? x : y;}
template<typename T1, typename T2> T1 chkmax (T1 x, T2 y) {return x < y ? y : x;}
#define si(x) (int)(x.size())
using namespace std;
typedef pair<int, int> PII;
typedef long long ll;
const int N = 1e5 + 10;
int n, a[N];
vector<int> e[N];
int sz[N], h[N], fa[N], son[N], dfn[N], idx, hs[N];
int pos[N], L[N], R[N];
void dfs1(int u, int father, int dep)
{
sz[u] = 1, h[u] = dep, fa[u] = father;
for (auto v : e[u])
{
if (v == father) continue;
dfs1(v, u, dep + 1);
sz[u] += sz[v];
if (sz[son[u]] < sz[v]) son[u] = v;
}
}
void dfs2(int u, int t)
{
dfn[u] = ++ idx, hs[idx] = u;
if (!son[u]) return;
dfs2(son[u], t);
for (auto v : e[u])
{
if (v == son[u] || v == fa[u]) continue;
dfs2(v, v);
}
}
struct E
{
int u, v;
}edge[N];
struct Ques
{
int l, r, id;
}q[N];
bool cmp(const Ques &a, const Ques &b)
{
if (pos[a.l] != pos[b.l]) return pos[a.l] < pos[b.l];
else return a.r < b.r;
}
int tot_c[N], cnt[N];
int res = 0;
void add(int p)
{
int x = a[hs[p]];
// cout << p << " " << hs[p] << endl;
res -= (tot_c[x] - cnt[x]) * cnt[x];
// cout << hs[p] << " " << x << endl;
cnt[x] ++;
res += (tot_c[x] - cnt[x]) * cnt[x];
// cout << p << " " << x << " " << res << endl;
}
void del(int p)
{
int x = a[hs[p]];
res -= (tot_c[x] - cnt[x]) * cnt[x];
// cout << p << " " << tot_c[x] - cnt[x] << endl;
cnt[x] --;
res += (tot_c[x] - cnt[x]) * cnt[x];
}
int ans[N];
signed main()
{
#ifdef DEBUG
freopen("in.in", "r", stdin);
freopen("out.out", "w", stdout);
#endif
ios::sync_with_stdio(false); cin.tie(nullptr), cout.tie(nullptr);
cin >> n;
for (int i = 1; i <= n; i ++)
{
cin >> a[i];
tot_c[a[i]] ++;
}
int len = sqrt(n);
for (int i = 1; i <= len; i ++)
{
L[i] = (i - 1) * len + 1;
R[i] = i * len;
}
if (R[len] < n) len ++, L[len] = R[len - 1] + 1, R[len] = n;
for (int i = 1; i <= len; i ++)
for (int j = L[i]; j <= R[i]; j ++)
pos[j] = i;
for (int i = 1; i < n; i ++)
{
int u, v; cin >> u >> v;
e[u].push_back(v);
e[v].push_back(u);
edge[i] = {u, v};
}
dfs1(1, 0, 1);
dfs2(1, 1);
// for (int i = 1; i <= n; i ++) cout << hs[i] << endl;
for (int i = 1; i < n; i ++)
{
auto [u, v] = edge[i];
if (h[u] < h[v]) swap(u, v);
// cout << "l r = " << dfn[u] << " " << dfn[u] + sz[u] - 1 << endl;
q[i] = {dfn[u], dfn[u] + sz[u] - 1, i};
}
sort(q + 1, q + n, cmp);
int ql = 1, qr = 0;
for (int i = 1; i < n; i ++)
{
auto [l, r, id] = q[i];
//cout << l << " " << r << " " << id << endl;
while (qr < r) add(++ qr);
// cout << id << " " << res << endl;
while (qr > r) del(qr --);
while (ql < l) del(ql ++);
while (ql > l) add(-- ql);
// cout << res << endl;
ans[id] = res;
}
for (int i = 1; i < n; i ++) cout << ans[i] << " ";
return 0;
}
AC了,大佬牛逼
哈哈,莫队还是可以看看,就相当于暴力优化for循环,但是可能正解不是莫队,你可以问问正解是啥,然后告诉我
我有个同学想做一下,请问一下能不能把链接发出来
https://pintia.cn/problem-sets/dashboard 邀请码:ed3ef8d9130e08a5
学校校赛的题目
我看了一眼其他同学AC的代码,都很一致的有add 和del 这两个操作函数,我想正解应该也是莫队吧
okok,谢谢啦
现在的初高中生都这么卷了,大学生该何去何从