提出 #12162502

---

ソースコード 拡げる

#include <iostream>
#include <vector>
#include <limits.h>
#include <algorithm>
#include <string>
#include <math.h>
#include <limits.h>
#include <queue>
#include <map>
#include <set>
#include <iomanip>
#include <bitset>
#include <cassert>
#include <random>
#include <functional>
#include <stack>
#include <iomanip>
#include <cassert>
//#include <boost/multiprecision/cpp_int.hpp>
#include <complex>
#include <cstdio>
#include <list>
#include <bitset>

//< in.txt > out.txt
using namespace std;
//std::ios::sync_with_stdio(false);
//std::cin.tie(0);
const long long MOD = 1e9+7;
const long long INF = 1e18;
typedef long long LL;
typedef long double LD;
typedef pair<LL, LL> PLL;
typedef pair<LD, LL> pdl;
typedef pair<LD, LD> pdd;
typedef vector<LL> VLL;
typedef vector<VLL> VVLL;
typedef unsigned long long ULL;
//typedef boost::multiprecision::cpp_int bigint;

template<class T>
struct Gedge {
	LL src, to;
	T cost;
	Gedge(LL s, LL t, T c) :src(s), to(t), cost(c) {}
	Gedge(LL t, T c) :src(-1), to(t), cost(c) {}
};

template<class T>
using WeightedGraph = vector<vector<Gedge<T>>>;
using UnweightedGraph = vector<vector<LL>>;
template<class T>
using Gedges = vector<Gedge<T>>;

struct TNode {
	LL parent;
	VLL childs;
	TNode() :parent(-1) {};
};
using Tree = vector<TNode>;

void SetTree(UnweightedGraph& G, Tree& T, LL root = 0) {
	T.resize(G.size());
	stack<PLL> st;
	st.push(PLL(root, -1));
	while (!st.empty()) {
		LL cur = st.top().first;
		LL p = st.top().second;
		st.pop();
		T[cur].parent = p;
		for (LL ch : G[cur]) {
			if (ch == p)continue;
			T[cur].childs.push_back(ch);
			st.push(PLL(ch, cur));
		}
	}
}

VLL ans;
LL N;
Tree tall;
VLL colors;
vector<Tree> tcol;
VLL tcindex;   //各頂点がその色ごとの木で何番目の頂点に位置するか
VVLL cind;   //各色木の各頂点と、元頂点の対応表
vector<stack<LL>> bfsst;
VLL treesize;

void bfs(LL n) {
	LL mycol = colors[n];
	tcindex[n] = tcol[mycol].size();
	tcol[mycol].push_back(TNode());
	LL scparent = bfsst[mycol].top();
	tcol[mycol].back().parent = bfsst[mycol].top();
	LL pind = tcindex[scparent];
	tcol[mycol][pind].childs.push_back(tcindex[n]);
	bfsst[mycol].push(n);
	treesize[n] = 1;
	cind[mycol].push_back(n);
	LL chlook = 0;
	for (LL ch : tall[n].childs) {
		bfs(ch);
		LL temp = treesize[ch];
		for (; chlook < tcol[mycol][tcindex[n]].childs.size(); chlook++) {
			LL chind = tcol[mycol][tcindex[n]].childs[chlook];
			temp -= treesize[cind[mycol][chind]];
		}
		ans[mycol] += temp*temp;
		treesize[n] += treesize[ch];
	}
	LL temp = treesize[n] - 1;
	for (LL ch : tcol[mycol][tcindex[n]].childs) {
		LL ordind = cind[mycol][ch];
		temp -= treesize[ordind];
	}
	//ans[mycol] += temp * temp;
	bfsst[mycol].pop();
}

void bfs2(LL n) {
	
}

int main() {
	std::ios::sync_with_stdio(false);
	std::cin.tie(0);
	cin >> N;
	colors.resize(N+1,-1);
	for (LL n = 1; n <= N; n++) {
		cin >> colors[n];
		colors[n]--;
	}
	{
		UnweightedGraph G(N+1);
		G[0].push_back(1);
		G[1].push_back(0);
		for (LL n = 0; n < N-1; n++) {
			LL a, b;
			cin >> a >> b;
			G[a].push_back(b);
			G[b].push_back(a);
		}
		SetTree(G, tall);
	}
	tcol.resize(N);
	bfsst.resize(N);
	tcindex.resize(N+1);
	cind.resize(N);
	for (LL n = 0; n < N; n++) {
		tcol[n].push_back(TNode());
		bfsst[n].push(0);
		cind[n].push_back(0);
	}
	treesize.resize(N + 1);
	ans.resize(N,0);
	cind.resize(N);
	bfs(1);
	for (LL col = 0; col < N; col++) {
		LL temp = N;
		for (LL ch : tcol[col][0].childs) {
			temp -= treesize[cind[col][ch]];
		}
		ans[col] += temp * temp;
	}
	for (LL col = 0; col < N; col++) {
		LL temp = (ans[col] - N - 1 + tcol[col].size()) / 2 + N + 1 - tcol[col].size();
		cout << (N + 1) * N / 2 - temp << "\n";
		//cout << temp << "\n\n";
	}
	return 0;
}

提出情報

コンパイルエラー

./Main.cpp:29:17: warning: unused variable 'MOD' [-Wunused-const-variable]
const long long MOD = 1e9+7;
                ^
./Main.cpp:30:17: warning: unused variable 'INF' [-Wunused-const-variable]
const long long INF = 1e18;
                ^
2 warnings generated.

ジャッジ結果