PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: pols_agyi_pols
Tester & Editorialist: iceknight1093
DIFFICULTY:
2688
PREREQUISITES:
DFS and subtree dp, combinatorics
PROBLEM:
You’re given a tree with N vertices, representing N cities and connections between them.
City i needs A_i towers built in it.
Initially, city 1 has a brick factory, and no other cities do. You can perform the following moves:
- Build a brick factory at some city, such that at least one neighbor of this city already has a brick factory.
- Build one tower at a city with a brick factory.
Find the number of ways of building all the towers using the minimum number of moves.
EXPLANATION:
Notice that the condition on building brick factors essentially boils down to:
“Root the tree at vertex 1. Then, a brick factory can be built at vertex u only if a factory has already been built at the parent of u”.
Let’s deal with a slightly simpler problem first: suppose we didn’t have to build any towers, and just wanted to build one brick factory at every city (still following the above rule).
How many possible ways are there of doing this?
Answer
This is a somewhat well-known idea, and can be found with the help of dynamic programming.
Let dp_u be the number of ways to perform the process for vertices in the subtree of u.
Let sz_u denote the size of the subtree of u.
Then, for any u whose children are v_1, v_2, \ldots, v_k, we have:
- Exactly sz_u moves will be made in this subtree.
Of them, the first is of course to build a factory at u itself. - After that, note that the process independently continues to each of the children of u.
We can perform moves in the subtree of each child, and moves performed within one subtree don’t affect moves performed in other subtrees: they can be interspersed as we like. - So, we can do the following:
- From the remaining sz_u - 1 positions, sz_{v_1} of them go to the first subtree, sz_{v_2} go to the second subtree, and so on.
There are
\displaystyle\binom{sz_u-1}{sz_{v_1}} \binom{sz_u-1-sz_{v_1}}{sz_{v_2}} \ldots \binom{sz_u-1-sz_{v_1}-\ldots - sz_{v_{k-1}}}{sz_{v_k}}
ways to perform this distribution, which can be further simplified down to
\displaystyle \frac{(sz_u-1)!}{(sz_{v_1})! (sz_{v_2})! \ldots (sz_{v_k})!} - After distributing positions, each subtree has its own orders as well.
This gives us a further multiplier of dp_{v_1}\cdot dp_{v_2} \cdot \ldots \cdot dp_{v_k}.
- From the remaining sz_u - 1 positions, sz_{v_1} of them go to the first subtree, sz_{v_2} go to the second subtree, and so on.
- With the base condition dp_u=1 for a leaf, this allows us to compute all the dp_u values; and hence find dp_1 which is what we wanted.
We’ll adapt this method to our current problem as well.
Let’s first find the minimum number of moves needed.
Clearly, we need one move for each tower; so that’s A_1 + A_2 + \ldots + A_N moves.
Apart from that, every city with A_i \gt 0 needs a brick factory built in it.
This further means that if A_u \gt 0, then every ancestor of u needs a brick factory built.
Turning this around, we see that city u needs a brick factory built if and only if some city in the subtree of u requires a tower; and subtree sums can be precomputed for all vertices in \mathcal{O}(N) using a dfs.
Now we need to count the number of orders of performing these moves.
That can be done using subtree dp in a similar fashion as what was described above.
Let \text{moves}[u] denote the number of moves we need for the subtree rooted at u, and \text{dp}[u] denote the number of ways to perform these moves.
Then,
- As noted above, \text{moves}[u] = A_u + \sum \text{moves}[v] across all children v of u.
Further, if \text{moves}[u] \gt 0, further add 1 to it, because we need to build a brick factory at u first. - Next, we’ll compute \text{dp}[u].
The logic is pretty much the exact same as what was done for the easier version:- The first move is fixed, we build a brick factory at u.
- After that, moves within its children are independent; and these are also independent moves from building towers at u.
- So, of the remaining \text{moves}[u]-1 positions:
- Fix A_u of them to build towers at for city u.
- For each child v of u, fix \text{moves}[v] positions where the operations on this subtree will occur, and also account for all orders of moves within this subtree itself.
- Together, this (upon a bit of reduction) gives us:
where the products are across all children v of u.
If factorials and inverse factorials are precomputed, this solves the problem in \mathcal{O}(N) time: the final answer is just \text{dp}[1].
Note that \text{moves}[1] can be as large as 10^6 + N, so make sure to precompute upto an appropriate upper bound (and not just till 10^6).
As an aside, the first ‘simpler’ problem we solved in fact somewhat directly solves this version too!
Consider a new tree where for each u, we attach a chain of size A_u to it as a subtree.
The answer is then just the number of orderings of this new tree (after pruning useless subtrees), which we already saw how to compute.
Of course, directly implementing this will result in a TLE because it’s proportional to the sum of A_i, and that’s bounded for a single testcase but not over all testcases.
However, it does explain why the final expressions end up looking remarkably similar.
TIME COMPLEXITY
Precomputation of 2\cdot 10^6 values, followed by \mathcal{O}(N) per testcase.
CODE:
Author's code (C++)
#include <bits/stdc++.h>
using namespace std;
#define ll long long
ll Mod=1000000007;
vector <ll> v[200005];
ll tower[200005];
const ll N=2*1e6+1;
ll factorialNumInverse[N + 1];
ll naturalNumInverse[N + 1];
ll fact[N + 1];
void InverseofNumber(ll p)
{
naturalNumInverse[0] = naturalNumInverse[1] = 1;
for (int i = 2; i <= N; i++)
naturalNumInverse[i] = naturalNumInverse[p % i] * (p - p / i) % p;
}
void InverseofFactorial(ll p)
{
factorialNumInverse[0] = factorialNumInverse[1] = 1;
for (int i = 2; i <= N; i++)
factorialNumInverse[i] = (naturalNumInverse[i] * factorialNumInverse[i - 1]) % p;
}
void factorial(ll p)
{
fact[0] = 1;
for (int i = 1; i <= N; i++) {
fact[i] = (fact[i - 1] * i) % p;
}
}
ll Binomial(ll N, ll R, ll p)
{
ll ans = ((fact[N] * factorialNumInverse[R])
% p * factorialNumInverse[N - R])
% p;
return ans;
}
void start()
{
ll p = 1000000007;
InverseofNumber(p);
InverseofFactorial(p);
factorial(p);
}
pair <ll,ll> dfs(ll pos,ll par){
ll r=tower[pos];
ll ans=1;
ll cnt=0;
pair <ll,ll> pp;
for(auto it:v[pos]){
if(it!=par){
pp=dfs(it,pos);
if(pp.second>0){
ans*=pp.first;
ans%=Mod;
ans*=factorialNumInverse[pp.second];
ans%=Mod;
cnt+=pp.second;
}
}
}
ans*=fact[cnt];
ans%=Mod;
cnt+=r;
ans*=Binomial(cnt,r,Mod);
ans%=Mod;
if(cnt){
cnt++;
}
return {ans,cnt};
}
int main(){
ios_base::sync_with_stdio(false);
cin.tie(NULL);
ll kitne_cases_hain;
kitne_cases_hain=1;
//freopen("input.txt","r",stdin);freopen("output.txt","w",stdout);
cin>>kitne_cases_hain;
start();
while(kitne_cases_hain--){
ll n;
cin>>n;
for(int i=1;i<=n;i++){
v[i].clear();
}
for(int i=1;i<=n;i++){
cin>>tower[i];
}
ll x,y;
for(int i=1;i<n;i++){
cin>>x>>y;
v[x].push_back(y);
v[y].push_back(x);
}
pair <ll,ll> p=dfs(1,0);
ll ans=p.first;
cout<<ans<<"\n";
}
return 0;
}
Editorialist's code (Python)
mod = 10**9 + 7
maxn = 10**6 + (2 * 10**5)
fac = [1] * maxn
ifac = [0] * maxn
for i in range(2, maxn): fac[i] = i * fac[i-1] % mod
ifac[-1] = pow(fac[-1], mod-2, mod)
for i in reversed(range(maxn-1)): ifac[i] = ifac[i+1] * (i+1) % mod
def C(n, r):
if n < r or r < 0: return 0
return fac[n] * ifac[r] % mod * ifac[n-r] % mod
for _ in range(int(input())):
n = int(input())
a = list(map(int, input().split()))
graph = [ [] for _ in range(n)]
for i in range(n-1):
u, v = map(int, input().split())
graph[u-1].append(v-1)
graph[v-1].append(u-1)
used, par, queue = [False]*n, [-1]*n, [0]
for x in queue:
used[x] = True
for v in graph[x]:
if not used[v]:
par[v] = x
queue.append(v)
moves, ways = [0]*n, [1]*n
for x in reversed(queue):
moves[x] += a[x]
rem = moves[x]
for v in graph[x]:
if v == par[x]: continue
ways[x] = ways[x] * C(rem, moves[v]) % mod * ways[v] % mod
rem -= moves[v]
if x > 0 and moves[x] > 0: moves[x] += 1
if x > 0: moves[par[x]] += moves[x]
print(ways[0])
