[toph.co] Sabiha and the War Victims

Author            : Dipu Kumar Mohanto 
                    CSE, Batch - 6
                    BRUR.
Problem Statement : Sabiha and the War Victims
Source            : toph.co
Category          : Graph Theory
Algorithm         : DSU on Tree
Verdict           : Accepted

#include "bits/stdc++.h"
 
using namespace std;
 
#define ll                 long long
#define MOD                1000000007
 
static const int MAXN = 2e5 + 5;
 
vector <int> graph[MAXN];
bool big[MAXN];
int sz[MAXN], tshirt[MAXN], temp[MAXN], cnt[MAXN];
int odd;
ll fact[MAXN], ifact[MAXN], up, down;
ll ans[MAXN];
 
map <int, int> mapper;
 
ll modPow(ll a, ll b, ll m)
{
    if (b == 1) return a % m;
    if (b & 1) return (a % m * modPow(a, b-1, m) % m) % m;
    ll ret = modPow(a, b>>1, m);
    return (ret % m * ret % m) % m;
}
 
ll modInv(ll a, ll m)
{
    return modPow(a, m-2, m);
}
 
void factorial()
{
    fact[0] = 1;
    ifact[0] = 1;
    for (int i = 1; i < MAXN; i++)
    {
        fact[i] = (i * fact[i-1]) %MOD;
        ifact[i] = modInv(fact[i], MOD);
    }
}
 
void dfs1(int u, int p=-1)
{
    sz[u] = 1;
    int len = graph[u].size();
    for (int i = 0; i < len; i++)
    {
        int v = graph[u][i];
        if (v == p) continue;
        dfs1(v, u);
        sz[u] += sz[v];
    }
}
 
void add(int u, int p, int x)
{
    int c = tshirt[u];
    !(cnt[c] & 1) ? odd++ : odd--;
    up -= cnt[c];
    down = (down * ifact[cnt[c]>>1]) %MOD;
    cnt[c] += x;
    up += cnt[c];
    down = (down * fact[cnt[c]>>1]) %MOD;
    int len = graph[u].size();
    for (int i = 0; i < len; i++)
    {
        int v = graph[u][i];
        if (v == p || big[v]) continue;
        add(v, u, x);
    }
}
 
void remov(int u, int p, int x)
{
    int c = tshirt[u];
    cnt[c] & 1 ? odd-- : odd++;
    up -= cnt[c];
    down = (down * ifact[cnt[c]>>1]) %MOD;
    cnt[c] -= x;
    up += cnt[c];
    down = (down * fact[cnt[c]>>1]) %MOD;
    int len = graph[u].size();
    for (int i = 0; i < len; i++)
    {
        int v = graph[u][i];
        if (v == p || big[v]) continue;
        remov(v, u, x);
    }
}
 
void dfs(int u, int p, int keep)
{
    int mx(-1), bigChild(-1);
    int len = graph[u].size();
    for (int i = 0; i < len; i++)
    {
        int v = graph[u][i];
        if (v == p) continue;
        if (sz[v] > mx) mx = sz[v], bigChild = v;
    }
    for (int i = 0; i < len; i++)
    {
        int v = graph[u][i];
        if (v == p || bigChild == v) continue;
        dfs(v, u, 0);
    }
    if (bigChild != -1) dfs(bigChild, u, 1), big[bigChild] = 1;
    add(u, p, 1);
    ll res = 0;
    if (odd <= 1)
    {
        res = (fact[up>>1] * modInv(down, MOD)) %MOD;
    }
    ans[u] = res;
    if (bigChild != -1) big[bigChild] = 0;
    if (keep == 0) remov(u, p, 1);
}
 
int main()
{
    //freopen("in.txt", "r", stdin);
 
    factorial();
    up = 0;
    down = 1;
 
    int tNode;
    scanf("%d", &tNode);
    for (int i = 1; i <= tNode; i++)
    {
        scanf("%d", tshirt+i);
        temp[i] = tshirt[i];
    }
    sort(temp+1, temp+tNode+1);
    int id = 0;
    for (int i = 1; i <= tNode; i++)
    {
        int x = temp[i];
        if (mapper.count(x) == 0)
        {
            id++;
            mapper[x] = id;
        }
    }
    for (int i = 1; i <= tNode; i++)
    {
        tshirt[i] = mapper[tshirt[i]];
    }
    for (int edge = 1; edge < tNode; edge++)
    {
        int u, v;
        scanf("%d %d", &u, &v);
        graph[u].push_back(v);
        graph[v].push_back(u);
    }
    dfs1(1);
    dfs(1, -1, 0);
    for (int i = 1; i <= tNode; i++)
    {
        printf("%lld", ans[i]);
        printf(i==tNode ? "\n" : " ");
    }
}