PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: gunpoint_88
Tester: iceknight1093
Editorialist: iceknight1093
DIFFICULTY:
TBD
PREREQUISITES:
DFS
PROBLEM:
You’re given a tree with N vertices and an array D of size N.
You must assign the values in D to the vertices of the tree.
A game map is a pair (X, Y) (1 \leq X \leq N, 0 \leq Y \lt N) denoting a choice of all vertices at distance at most Y from X.
The difficulty of a game map is the sum of all vertices included in it.
You want to minimize the difficulty.
Under an optimal assignment of values of D to the vertices, find the minimum expected difficulty if any one of the game maps can be chosen with equal probability.
EXPLANATION:
Let’s first learn how to compute the expected difficulty of a game map, given an assignment of the D values to the vertices.
Let A_i denote the value assigned to vertex i.
Each map is equally likely to be chosen, and there are N^2 of them.
So, we can instead find the total sum of difficulties of all maps, and then divide this by N^2.
Finding this total sum can be done quickly with the technique of contribution, that is, look at how much individual part contributes to the whole.
Consider some vertex u — let’s count the number of maps that include A_u in their sum.
Suppose a game map (X, Y) includes A_u.
This means \text{dist}(X, u) \leq Y, by definition.
Turning this around, if d = \text{dist}(X, u), A_u will appear in each of the maps (X, d), (X, d+1), (X, d+2), \ldots, (X, N-1).
This is a total of N-d maps.
So, overall, A_u appears in \sum_{X=1}^N (N - \text{dist}(X, u)) maps.
Define S_u = \sum_{v=1}^N \text{dist}(u, v).
The total sum is thus
\sum_{u=1}^N A_u\cdot \left( N^2 - S_u\right) = N^2\sum_{u=1}^N A_u - \sum_{u=1}^N A_u\cdot S_u
The first term is a constant, since the sum of all A_u is just the sum of all D_u.
So, to minimize this, we must focus on the second term, i.e, maximizing \sum A_uS_u.
There are two things to do here:
- Finding all the S_u values; and
- Figuring out an optimal assignment based on these values.
The first part is a rather classical application of the rerooting technique.
Notice that S_u is the sum of all distances of other vertices from u.
This can be computed with two DFS-es:
- First, DFS once to compute the sum of distances from u to vertices only within its subtree.
This should be straightforward to figure out, you’ll likely need to maintain subtree sizes as well. - Then, perform a second DFS to account for vertices outside the subtree of each vertex.
During this DFS, you pass on the results from parent to child.
If this type of algorithm is new to you, an example can be found in the link just above.
Now, suppose we know all the S_u values.
We need to match the S_u values with the values in D so as to maximize the sum of their pointwise products.
This can be done greedily: sort S and D in ascending order, then match S_i with D_i.
This is optimal due to the rearrangement inequality.
Remember to divide by N^2 in the end, since we want the expected value.
Division under modulo requires modular inverses.
TIME COMPLEXITY:
\mathcal{O}(N\log N) per testcase.
CODE:
Author's code (C++)
#include<bits/stdc++.h>
using namespace std;
using ll=long long;
const ll mod=1e9+7;
#ifdef ANI
#include "D:/DUSTBIN/local_inc.h"
#else
#define dbg(...) 0
#endif
/*
generation plan:
1 <= T <= 1e4
1 <= N <= 3e5
1 <= Di <= 1e9
file 1:
all non-isomorphic trees of size <= 16
(there are 19320 of these as per oeis, sum of n over all these is < 3e5)
file 2 to 5:
randomly generated trees with N_max varying as:
file 2: 50
file 3: 500
file 4: 5000
file 5: 50000
file 6 to 8:
Instead of Files 6 to 8 being totally random, I'd prefer something like this:
6: Every node has either 2 or 3 children
7: n/2 length chain, and every node has one other child leaf (maybe add some small perturbations to this)
8: Star-ish graph. sqrt(n) children of the root. Each child having [sqrt(n) - eps, sqrt(n) + eps] length chain beneath it
file 9:
N=max
line graph
file 10 to 12:
N=max
sparse graphs (random generation creates dense graphs)
*/
class testcase{
public:
ll n;
vector<array<ll,2>> tree;
vector<vector<ll>> e;
vector<ll> d;
ll ans;
testcase(vector<array<ll,2>> tree,vector<ll> d) {
this->n=d.size();
this->tree=tree;
this->d=d;
e=vector<vector<ll>>(n);
for(auto el:tree) {
assert(el[0]<n && el[1]<n && el[0]>-1 && el[1]>-1);
e[el[0]].push_back(el[1]);
e[el[1]].push_back(el[0]);
}
this->ans=solution();
}
testcase(vector<vector<ll>> e,vector<ll> d) {
this->n=d.size();
this->e=e;
this->d=d;
tree=vector<array<ll,2>>();
for(ll i=0;i<n;i++) {
for(ll j:e[i]) {
assert(j<n && j>=0);
if(i<j) tree.push_back({i,j});
}
}
this->ans=solution();
assert(this->ans>=0);
}
// solution
ll fastexp(ll base,ll exp) {
ll res=1;
base%=mod;
while(exp) {
if(exp&1) res=(res*base)%mod;
base=(base*base)%mod;
exp>>=1;
}
return res;
}
ll modinv(ll num) {
return fastexp(num, mod-2);
}
ll solution() {
vector<ll> sub(n,0),ct(n,1);
function<void(ll,ll)> dfs=[&](ll cur,ll par)->void{
for(ll node:e[cur]) {
if(node^par) {
dfs(node,cur);
sub[cur]+=sub[node]+ct[node];
ct[cur]+=ct[node];
}
}
};
dfs(0,-1);
vector<ll> tot(n,0);
function<void(ll,ll,ll)> dfs2=[&](ll cur,ll par,ll anc)->void{
tot[cur]=sub[cur]+(anc+(n-ct[cur]));
for(ll node:e[cur]) {
if(node^par) {
dfs2(node,cur,tot[cur]-sub[node]-ct[node]);
}
}
};
dfs2(0,-1,0);
ll ans=0;
sort(tot.begin(),tot.end());
sort(d.begin(),d.end());
for(ll i=0;i<n;i++) {
ans+=((n*n-tot[i])%mod)*d[i]%mod;
ans%=mod;
}
return (ans*modinv(n*n)%mod);
}
void write(ofstream &inp, ofstream &oup) {
inp<<n<<"\n";
for(ll i=0;i<n;i++) {
inp<<d[i]<<" \n"[i==n-1];
}
for(ll i=0;i<n-1;i++) {
inp<<tree[i][0]+1<<" "<<tree[i][1]+1<<"\n";
}
oup<<ans<<"\n";
}
};
namespace isotree{
const ll mod=29996224275833;
vector<ll> is_prime,primes;
void sieve(ll n=1e6) {
is_prime=vector<ll>(n+1,1);
is_prime[0]=is_prime[1]=0;
for(ll i=2;i*i<=n;i++) {
if(!is_prime[i]) continue;
for(ll j=i*i;j<=n;j+=i) {
is_prime[j]=0;
}
}
for(ll i=2;i<=n;i++)
if(is_prime[i]) primes.push_back(i);
}
ll fastexp(ll base,ll exp,ll mod) {
ll res=1;
base%=mod;
while(exp) {
if(exp&1) res=(res*base)%mod;
base=(base*base)%mod;
exp>>=1;
}
return res;
}
ll modinv(ll num,ll mod) {
return fastexp(num, mod-2,mod);
}
class tree_hasher{
public:
ll n,mod,seed,root,p;
vector<vector<ll>> e;
vector<ll> subtree_hashes,rooted_hashes;
ll isomorphic_hash=0;
tree_hasher(vector<vector<ll>> e,ll mod=1e9+7,ll seed=2746263,ll root=0) {
this->mod=mod;
this->e=e;
this->seed=seed;
this->n=e.size();
this->root=root;
subtree_hashes=rooted_hashes=vector<ll>(n,0);
compute_subtree(root,-1);
compute_rooted(root,-1,-1);
compute_isomorphic_hash();
}
void compute_isomorphic_hash() {
vector<ll> d(n,1),centroids;
function<void(ll,ll)> dfs=[&](ll cur,ll par)->void{
ll tot_ct=0,mx_subt=0;
for(ll node:e[cur]) {
if(node==par) continue;
dfs(node,cur);
d[cur]+=d[node];
tot_ct+=d[node];
mx_subt=max(mx_subt,d[node]);
}
mx_subt=max(mx_subt,n-1-tot_ct);
if(2*mx_subt<=n) {
centroids.push_back(cur);
}
};
dfs(0,-1);
assert(centroids.size()<=2 && centroids.size()>=1);
array<ll,2> is_hsh={0,0};
is_hsh[0]=rooted_hashes[centroids[0]];
if(centroids.size()==2) {
is_hsh[1]=rooted_hashes[centroids[1]];
}
if(is_hsh[1]>is_hsh[0]) swap(is_hsh[0],is_hsh[1]);
isomorphic_hash=((is_hsh[0]>>5)^(is_hsh[1]<<6)^(is_hsh[0]*seed%mod)^(is_hsh[1]) + seed)%mod;
}
void compute_subtree(ll cur,ll par) {
vector<ll> c;
for(ll node:e[cur]) {
if(node==par) continue;
compute_subtree(node,cur);
c.push_back(subtree_hashes[node]);
}
sort(c.begin(),c.end());
ll res=seed%mod,k=c.size();
ll p=*upper_bound(primes.begin(),primes.end(),k);
for(ll i=0;i<c.size();i++) {
res+=(c[i]*c[i]%mod + seed + c[i]*fastexp(p,i,mod)%mod)%mod;
res%=mod;
}
subtree_hashes[cur]=res;
}
void compute_rooted(ll cur,ll par,ll anc_h) {
vector<pair<ll,ll>> h;
if(par!=-1) h.push_back({anc_h,-1});
for(ll node:e[cur]) {
if(node==par) continue;
h.push_back({subtree_hashes[node],node});
}
sort(h.begin(),h.end());
map<ll,ll> idx; // index of this particular node in sorted h
ll k=h.size();
vector<ll> pref_h(k+1,0),pref_c(k+1,0);
ll p=*upper_bound(primes.begin(),primes.end(),k-1);
ll pp=*upper_bound(primes.begin(),primes.end(),k);
ll ans=seed;
for(ll i=0;i<k;i++) {
ll node=h[i].second,hsh=h[i].first;
idx[node]=i;
pref_h[i+1]=(pref_h[i] + hsh*hsh%mod + seed + hsh*fastexp(p,i,mod)%mod)%mod;
pref_c[i+1]=(pref_c[i] + hsh*hsh + seed)%mod;
ans+=(pref_h[i] + hsh*hsh%mod + seed + hsh*fastexp(pp,i,mod)%mod)%mod;
}
rooted_hashes[cur]=ans;
for(ll node:e[cur]) {
if(node==par) continue;
ll ii=idx[node];
/*
H: [h1 h2 h3 h4] (h5) [h6 h7 h8]
hash of [0...idx-1] merge with hash of [idx+1....k-1]
*/
ll h1=pref_h[ii];
ll h2=(pref_h[k]-pref_h[ii+1]+mod)%mod;
ll ext=(pref_c[k]-pref_c[ii+1]+mod)%mod;
h2=(h2-ext+mod)%mod;
h2*=modinv(p,mod);
h2=(h2+ext)%mod;
ll pass_hash=(h1+h2+seed)%mod;
compute_rooted(node,cur,pass_hash);
}
}
};
using tree=vector<vector<ll>>;
vector<tree> generate_all(ll n) {
vector<set<ll>> hashes(n+1);
vector<vector<tree>> trees(n+1);
tree base_tree={{}};
trees[1]={base_tree};
hashes[1].insert(tree_hasher(base_tree).isomorphic_hash);
cout<<"n = 1, no. of trees = 1"<<endl;
vector<tree> res={base_tree};
for(ll sz=2;sz<=n;sz++) {
ll new_node=sz-1;
for(auto x:trees[sz-1]) {
x.push_back({});
for(ll to=0;to<new_node;to++) {
// if(x[to].size()>=4) continue;
x[to].push_back(new_node);
x[new_node].push_back(to);
ll hsh=tree_hasher(x).isomorphic_hash*tree_hasher(x,998244353,737636,0).isomorphic_hash;
if(!hashes[sz].count(hsh)) {
hashes[sz].insert(hsh);
trees[sz].push_back(x);
}
x[to].pop_back();
x[new_node].pop_back();
}
}
for(auto x:trees[sz]) {
res.push_back(x);
}
cout<<"n = "<<sz<<", no. of trees = "<<hashes[sz].size()<<endl;
}
return res;
}
};
vector<ll> perm(ll n) {
vector<ll> res(n);
mt19937_64 mt(rand());
iota(res.begin(),res.end(),0);
for(ll i=n-1;i>0;i--) {
swap(res[i],res[mt()%i]);
}
return res;
}
vector<array<ll,2>> simple_tree(ll n) {
mt19937_64 mt(rand());
vector<array<ll,2>> res(n-1);
for(ll i=1;i<n;i++) {
res[i-1]={i,mt()%i};
}
return res;
}
vector<array<ll,2>> prufer_tree(ll n) {
mt19937_64 mt(rand());
vector<ll> prufer(n-2);
for(ll i=0;i<n-2;i++) {
prufer[i]=mt()%n;
}
vector<ll> deg(n,1);
for(ll i:prufer)
deg[i]++;
set<ll> leaves;
for(ll i=0;i<n;i++) {
if(deg[i]==1)
leaves.insert(i);
}
vector<array<ll,2>> res;
for(ll u:prufer) {
ll leaf=(*leaves.begin());
leaves.erase(leaves.begin());
res.push_back({leaf,u});
deg[u]-=1;
if(deg[u]==1) {
leaves.insert(u);
}
}
res.push_back({*leaves.begin(),n-1});
return res;
}
vector<array<ll,2>> child_ct_tree(ll n,ll L=2,ll R=2) {
mt19937_64 mt(rand());
if(n==3) {
return {{0,1},{0,2}};
}
vector<array<ll,2>> tree;
vector<ll> leaf={0},prev,next_leaf;
ll tot=n-1,cur=1;
auto add_leaf=[&](ll base)->void{
tree.push_back({base,cur});
next_leaf.push_back(cur++);
tot--;
};
auto rng=[&](ll l,ll r)->ll{
return l+mt()%(r-l+1);
};
for(ll level=1;cur<n;level++) {
ll k=leaf.size();
for(ll i=0;i<k;i++) {
ll give=min(n-cur,rng(L,R));
for(ll j=0;j<give;j++) {
add_leaf(leaf[i]);
}
}
prev=leaf;
leaf=next_leaf;
next_leaf.clear();
}
return tree;
}
void file1() {
/*
all non-isomorphic unrooted trees of sizes [1..15]
*/
string inp_path="./testcases/test1.in";
string out_path="./testcases/test1.out";
ofstream input(inp_path);
ofstream output(out_path);
mt19937_64 mt(rand());
auto trees=isotree::generate_all(15);
vector<testcase> tc;
for(auto el:trees) {
ll n=el.size();
ll dmax=2*n;
vector<ll> d(n);
for(ll i=0;i<n;i++) {
d[i]=mt()%dmax;
}
tc.push_back(testcase(el,d));
}
input<<tc.size()<<"\n";
for(auto &el:tc) {
el.write(input,output);
}
cout<<inp_path<<" finished"<<endl;
}
void generic(ll tn,ll nmax,ll dmax=1e9,ll fixed=0) {
string inp_path="./testcases/test"+to_string(tn)+".in";
string out_path="./testcases/test"+to_string(tn)+".out";
ofstream input(inp_path);
ofstream output(out_path);
mt19937_64 mt(rand());
vector<testcase> tc;
ll nsum=3e5, tsum=2e4;
while(nsum>=16 && tc.size()<tsum) {
ll n=1ll+mt()%min(nmax, nsum);
n=max(n,16ll);
nsum-=n;
auto tree=prufer_tree(n);
vector<ll> d(n);
for(ll i=0;i<n;i++) {
d[i]=mt()%dmax;
}
tc.push_back(testcase(tree,d));
}
input<<tc.size()<<"\n";
for(auto &el:tc) {
el.write(input, output);
}
cout<<inp_path<<" finished"<<endl;
}
void file6() {
string inp_path="./testcases/test6.in";
string out_path="./testcases/test6.out";
ofstream input(inp_path);
ofstream output(out_path);
mt19937_64 mt(rand());
ll n=3e5,dmax=1e9;
vector<ll> d(n);
for(ll i=0;i<n;i++) {
d[i]=mt()%dmax;
}
vector<testcase> tc={testcase(child_ct_tree(n,2,3),d)};
input<<tc.size()<<"\n";
for(auto &el:tc) {
el.write(input, output);
}
cout<<inp_path<<" finished"<<endl;
}
void file7() {
string inp_path="./testcases/test7.in";
string out_path="./testcases/test7.out";
ofstream input(inp_path);
ofstream output(out_path);
mt19937_64 mt(rand());
ll n=3e5-50,dmax=1e9;
vector<ll> d(n);
for(ll i=0;i<n;i++) {
d[i]=mt()%dmax;
}
vector<array<ll,2>> tree;
ll cur=0;
for(ll i=0;i<n/2;i++) {
tree.push_back({cur++,cur});
}
ll m=cur++;
for(ll i=1;i<m && cur<n; i++) {
if(mt()%9192) tree.push_back({i,cur++});
}
while(cur<n) {
ll to=mt()%cur;
tree.push_back({to,cur++});
}
vector<testcase> tc={testcase(tree,d)};
input<<tc.size()<<"\n";
for(auto &el:tc) {
el.write(input, output);
}
cout<<inp_path<<" finished"<<endl;
}
void file8() {
string inp_path="./testcases/test8.in";
string out_path="./testcases/test8.out";
ofstream input(inp_path);
ofstream output(out_path);
mt19937_64 mt(rand());
ll n=3e5-122,dmax=1e9;
vector<ll> d(n);
for(ll i=0;i<n;i++) {
d[i]=mt()%dmax;
}
vector<array<ll,2>> tree;
ll len=sqrt(n),cur=0;
for(ll i=0;i<len && cur<n;i++) {
tree.push_back({++cur,0});
for(ll j=0;j<len && cur<n;j++) {
tree.push_back({cur++,cur});
}
}
cur++;
while(cur<n) {
tree.push_back({mt()%cur,cur++});
}
vector<testcase> tc={testcase(tree,d)};
input<<tc.size()<<"\n";
for(auto &el:tc) {
el.write(input, output);
}
cout<<inp_path<<" finished"<<endl;
}
void file9() {
string inp_path="./testcases/test9.in";
string out_path="./testcases/test9.out";
ofstream input(inp_path);
ofstream output(out_path);
mt19937_64 mt(rand());
ll n=3e5-8,dmax=1e9;
vector<ll> d(n);
for(ll i=0;i<n;i++) {
d[i]=mt()%dmax;
}
vector<array<ll,2>> tree;
vector<ll> p=perm(n);
for(ll i=0;i<n-1;i++) {
tree.push_back({p[i],p[i+1]});
}
vector<testcase> tc={testcase(tree,d)};
input<<tc.size()<<"\n";
for(auto &el:tc) {
el.write(input, output);
}
cout<<inp_path<<" finished"<<endl;
}
void sparse(ll tn,ll k) {
string inp_path="./testcases/test"+to_string(tn)+".in";
string out_path="./testcases/test"+to_string(tn)+".out";
ofstream input(inp_path);
ofstream output(out_path);
mt19937_64 mt(rand());
ll n=3e5-8,dmax=1e9;
vector<ll> d(n);
for(ll i=0;i<n;i++) {
d[i]=mt()%dmax;
}
vector<array<ll,2>> tree;
for(ll i=0;i<k;i++) {
tree.push_back({i,i+1});
}
for(ll i=k+1;i<n;i++) {
tree.push_back({i,mt()%i});
}
vector<testcase> tc={testcase(tree,d)};
input<<tc.size()<<"\n";
for(auto &el:tc) {
el.write(input, output);
}
cout<<inp_path<<" finished"<<endl;
}
static void run_with_stack_size(void (*func)(void), size_t stsize) {
char *stack, *send;
stack = (char *)malloc(stsize);
send = stack + stsize - 16;
send = (char *)((uintptr_t)send / 16 * 16);
asm volatile(
"mov %%esp, (%0)\n"
"mov %0, %%esp\n"
:
: "r"(send));
func();
asm volatile("mov (%0), %%esp\n" : : "r"(send));
free(stack);
}
void main_() {
srand(1929138);
isotree::sieve();
file1();
generic(2,50);
generic(3,500);
generic(4,5000);
generic(5,50000);
file6();
file7();
file8();
file9();
sparse(10,3);
sparse(11,30);
sparse(12,12);
}
int main() {
// ll stck=1024 * 1024 * 1024;
// run_with_stack_size(main_, stck); // run with a 1 GiB stack
// return 0;
ll t; cin>>t;
while(t--) {
ll n; cin>>n;
vector<ll> d(n);
for(ll i=0;i<n;i++) {
cin>>d[i];
}
vector<array<ll,2>> tree;
for(ll i=0;i<n-1;i++) {
ll u,v; cin>>u>>v;
tree.push_back({u-1,v-1});
}
cout<<testcase(tree,d).ans<<"\n";
}
// auto t=child_ct_tree(6,3,3);
// dbg(t);
}
Editorialist's code (C++)
// #pragma GCC optimize("O3,unroll-loops")
// #pragma GCC target("avx2,bmi,bmi2,lzcnt,popcnt")
#include "bits/stdc++.h"
using namespace std;
using ll = long long int;
mt19937_64 rng(chrono::high_resolution_clock::now().time_since_epoch().count());
int main()
{
ios::sync_with_stdio(false); cin.tie(0);
int t; cin >> t;
while (t--) {
int N; cin >> N;
vector<int> D(N);
for (int &x : D) cin >> x;
sort(begin(D), end(D));
vector adj(N, vector<int>());
for (int i = 0; i < N-1; ++i) {
int u, v; cin >> u >> v;
adj[--u].push_back(--v);
adj[v].push_back(u);
}
vector<ll> distsum(N), subsz(N);
auto dfs = [&] (const auto &self, int u, int p) -> void {
subsz[u] = 1;
for (int v : adj[u]) {
if (v == p) continue;
self(self, v, u);
distsum[u] += distsum[v] + subsz[v];
subsz[u] += subsz[v];
}
};
auto reroot = [&] (const auto &self, int u, int p, ll up) -> void {
for (int v : adj[u]) {
if (v == p) continue;
self(self, v, u, up + distsum[u] - (distsum[v] + subsz[v]) + N - subsz[v]);
}
distsum[u] += up;
};
dfs(dfs, 0, 0);
reroot(reroot, 0, 0, 0);
sort(begin(distsum), end(distsum));
const int mod = 1e9 + 7;
ll ans = 0;
for (int i = 0; i < N; ++i) {
ans += D[i] % mod * ((1LL*N*N - distsum[i]) % mod);
ans %= mod;
}
int invN = 1;
{
int pw = mod - 2;
while (pw) {
if (pw & 1) invN = (1LL * invN * N) % mod;
N = (1LL * N * N) % mod;
pw /= 2;
}
}
ans = ((1LL * ans * invN) % mod * invN) % mod;
cout << ans << '\n';
}
}