Submission #4353103


Source Code Expand

Copy
#include <bits/stdc++.h>
#define rep(i,n) for(int i=0;i<(int)(n);i++)
#define rep1(i,n) for(int i=1;i<=(int)(n);i++)
#define all(c) c.begin(),c.end()
#define pb push_back
#define fs first
#define sc second
#define show(x) cout << #x << " = " << (x) << endl
#define chmin(x,y) x=min(x,y)
#define chmax(x,y) x=max(x,y)
using namespace std;
template<class S,class T> ostream& operator<<(ostream& o,const pair<S,T> &p){
	return o<<"("<<p.fs<<","<<p.sc<<")";
}
template<class T> ostream& operator<<(ostream& o,const vector<T> &vc){
	o<<"{";
	for(const T& v:vc) o<<v<<",";
	o<<"}";
	return o;
}
using ll = long long;
template<class T> using V = vector<T>;
template<class T> using VV = vector<vector<T>>;
int bsr(int x) { return 31 - __builtin_clz(x); }
using D = double;
const D pi = acos(-1);
using Pc = complex<D>;

void fft(bool type, vector<Pc> &c){	//multiply : false -> mult -> true
	static vector<Pc> buf[30];
	int N = c.size();
	int s = bsr(N);
	assert(1<<s == N);
	if(buf[s].empty()){
		buf[s]=vector<Pc>(N);
		rep(i,N) buf[s][i] = polar<D>(1,i*2*pi/N);
	}
	vector<Pc> a(N),b(N);
	copy(begin(c),end(c),begin(a));
	rep1(i,s){
		int W = 1<<(s-i);
		for(int y=0;y<N/2;y+=W){
			Pc now = buf[s][y];
			if(type) now = conj(now);
			rep(x,W){
				auto l =       a[y<<1 | x];
				auto r = now * a[y<<1 | x | W];
				b[y | x]        = l+r;
				b[y | x | N>>1] = l-r;
			}
		}
		swap(a,b);
	}
	copy(begin(a),end(a),begin(c));
}
template<class Mint>
vector<Mint> multiply_fft(const vector<Mint>& x,const vector<Mint>& y){
	if(x.empty() || y.empty()) return {};
	const int B = 15;
	const int K = 2;
	int S = x.size()+y.size()-1;
	int N = 1; while(N<S) N*=2;
	vector<Pc> a[K],b[K];
	rep(t,K){
		a[t] = vector<Pc>(N);
		b[t] = vector<Pc>(N);
		rep(i,x.size()) a[t][i] = Pc( (x[i] >> (t*B)) & ((1<<B)-1) , 0 );
		rep(i,y.size()) b[t][i] = Pc( (y[i] >> (t*B)) & ((1<<B)-1) , 0 );
		fft(false,a[t]);
		fft(false,b[t]);
	}
	vector<Mint> z(S);
	vector<Pc> c(N);
	Mint base = 1;
	rep(t,K+K-1){
		fill_n(begin(c),N,Pc(0,0));
		rep(xt,K){
			int yt = t-xt;
			if(0<=yt && yt<K){
				rep(i,N) c[i] += a[xt][i] * b[yt][i];
			}
		}
		fft(true,c);
		rep(i,S){
			c[i] *= 1.0/N;
			z[i] += Mint(ll(round(c[i].real()))) * base;
		}
		base *= 1<<B;
	}
	return z;
}

namespace Normalize{
	template<class E>
	void dfs(vector<vector<E>>& G,vector<int>& ord, int v,int p=-1){
		ord.pb(v);
		int K = G[v].size();
		rep(i,K){
			if(G[v][i].to==p){
				rotate(G[v].begin()+i,G[v].begin()+i+1,G[v].end());
				K--;
				break;
			}
		}
		rep(i,K){
			dfs(G,ord,G[v][i].to,v);
		}
	}
	template<class E>
	vector<int> normalize_and_gettopord(vector<vector<E>>& G, int r=0){
		vector<int> ord;
		dfs(G,ord,r);
		return ord;
	}
}

/*
	Node + に可換は仮定しなくても大丈夫(足す順序は変えてない(ただしupを足す場所は適当なので注意))
*/
template<class N>
struct BidirectionalTreeDP{
	vector<N> dp;	//dp[v] <- u1,u2,... 			pを削ってvが根
	vector<N> up;	//up[v] <- v's brothers + pp 	vを削ってpが根
	vector<N> rp;	//dp[r] <- rを根とした時のこたえ

	vector<int> par;

	BidirectionalTreeDP(){}

	template<class E>
	BidirectionalTreeDP(vector<vector<E>>& G, int r=0){
		int V = G.size();
		dp.assign(V,N());
		up.assign(V,N());
		rp.assign(V,N());
		par.assign(V,0);

		vector<int> ord = Normalize::normalize_and_gettopord<E>(G,r);
		rep(t,V){
			int v = ord[t];
			if(v==r) par[v]=-1;
			else par[v]=G[v].back().to;
		}

		for(int t=V-1;t>=0;t--){	//dfs
			int v = ord[t];
			dp[v] = N();
			int K = G[v].size() - (v!=r);
			rep(i,K){
				const E& e = G[v][i];
				int u = e.to;
				dp[v] = dp[v] + dp[u].append_edge(v,e);
			}
			dp[v].finalize(v);
		}

		rep(t,V){					//ufs
			int v = ord[t];
			int K = G[v].size() - (v!=r);
			vector<N> ls(K+1),rs(K+1);
			rep(i,K){
				ls[i+1] = ls[i] + dp[G[v][i].to].append_edge(v,G[v][i]);
				rs[K-1-i] = dp[G[v][K-1-i].to].append_edge(v,G[v][K-1-i]) + rs[K-i];
			}
			rep(i,K){
				const E& e = G[v][i];
				int u = e.to;
				up[u] = ls[i] + rs[i+1];
				if(v!=r) up[u] = up[u] + up[v].append_edge(v,G[v].back());
				up[u].finalize(v);
			}
			rp[v] = ls[K];
			if(v!=r) rp[v] = rp[v] + up[v].append_edge(v,G[v].back());
			rp[v].finalize(v);
		}
	}

	N get(int v,int p=-1){	//pを削ってvが根
		if(p==-1) return rp[v];
		if(par[v]==p) return dp[v];
		return up[p];
	}
};

/*
	get(v,p) = (部分木の直径, 部分木のmax dist, 2nd max dist)
*/

struct Node{	//
	int dia;
	array<int,2> rd;
	Node(){
		dia=0;
		rd[0]=rd[1]=0;
	}

	/*
		根付き木→森
		e=(p -> this)を追加したものを返す
	*/
	template<class E>
	Node append_edge(int p,const E& e) const {
		Node n;
		n.rd[0] = rd[0] + 1;
		n.rd[1] = 0;
		n.dia = dia;
		return n;
	}
	Node operator+(const Node& r) const {
		Node n;
		vector<int> vc;
		rep(t,2) vc.pb(rd[t]),vc.pb(r.rd[t]);
		sort(all(vc),greater<int>());
		rep(t,2) n.rd[t]=vc[t];
		n.dia = max(dia,r.dia);
		return n;
	}
	void finalize(int r){
		chmax(dia,rd[0]+rd[1]);
	}
};

struct Edge{
	int to;
};

int main(){
	cin.tie(0);
	ios::sync_with_stdio(false);		//DON'T USE scanf/printf/puts !!

	BidirectionalTreeDP<Node> A,B;
	int NA,NB;

	rep(t,2){
		int N;
		cin>>N;
		if(t==0) NA = N;
		if(t==1) NB = N;
		VV<Edge> G(N);
		rep(i,N-1){
			int x,y;
			cin>>x>>y;
			x--,y--;
			G[x].pb({y});
			G[y].pb({x});
		}
		(t==0 ? A : B) = BidirectionalTreeDP<Node>(G);
	}

	int old = max(A.get(0,-1).dia,B.get(0,-1).dia);

	V<ll> a(NA),b(NB+1);
	rep(i,NA){
		int x = A.get(i,-1).rd[0];
		a[x]++;
	}
	rep(i,NB){
		int x = B.get(i,-1).rd[0];
		b[x+1]++;
	}
	auto c = multiply_fft<ll>(a,b);

	ll ans = 0;
	rep(i,c.size()){
		ans += max(old,i) * c[i];
	}
	cout<<ans<<endl;
}

Submission Info

Submission Time
Task B - Bonsai Grafting
User sigma425
Language C++14 (GCC 5.4.1)
Score 700
Code Size 5931 Byte
Status
Exec Time 548 ms
Memory 50756 KB

Test Cases

Set Name Score / Max Score Test Cases
Sample 0 / 0 01.txt, 02.txt
All 700 / 700 01.txt, 02.txt, 11.txt, 12.txt, 13.txt, 14.txt, 15.txt, 16.txt, 17.txt, 18.txt, 19.txt, 20.txt, 21.txt, 22.txt, 23.txt, 24.txt, 25.txt, 26.txt, 27.txt, 28.txt, 29.txt, 30.txt, 31.txt, 32.txt, 33.txt, 34.txt, 35.txt, 36.txt, 37.txt, 38.txt, 39.txt, 40.txt, 41.txt, 42.txt, 43.txt, 44.txt, 45.txt, 46.txt, 47.txt, 48.txt, 49.txt, 50.txt
Case Name Status Exec Time Memory
01.txt 3 ms 384 KB
02.txt 1 ms 256 KB
11.txt 1 ms 256 KB
12.txt 1 ms 256 KB
13.txt 1 ms 256 KB
14.txt 1 ms 256 KB
15.txt 1 ms 256 KB
16.txt 1 ms 256 KB
17.txt 1 ms 256 KB
18.txt 1 ms 256 KB
19.txt 1 ms 256 KB
20.txt 1 ms 256 KB
21.txt 506 ms 46404 KB
22.txt 482 ms 44100 KB
23.txt 507 ms 46404 KB
24.txt 507 ms 49860 KB
25.txt 463 ms 44184 KB
26.txt 487 ms 46364 KB
27.txt 486 ms 49728 KB
28.txt 506 ms 46400 KB
29.txt 509 ms 49856 KB
30.txt 525 ms 50624 KB
31.txt 526 ms 46412 KB
32.txt 493 ms 44100 KB
33.txt 548 ms 46404 KB
34.txt 525 ms 50756 KB
35.txt 501 ms 44096 KB
36.txt 491 ms 46404 KB
37.txt 491 ms 48704 KB
38.txt 508 ms 46396 KB
39.txt 514 ms 49472 KB
40.txt 545 ms 49728 KB
41.txt 271 ms 22476 KB
42.txt 249 ms 22092 KB
43.txt 250 ms 22472 KB
44.txt 249 ms 22092 KB
45.txt 232 ms 22344 KB
46.txt 234 ms 22088 KB
47.txt 256 ms 26696 KB
48.txt 247 ms 22216 KB
49.txt 256 ms 24904 KB
50.txt 247 ms 24520 KB