BFSDFS - Editorial

PROBLEM LINK:

Contest Division 1
Contest Division 2
Contest Division 3
Contest Division 4

Setter: Abhinav Sharma
Testers: Jeevan Jyot Singh, Nishank Suresh
Editorialist: Abhinav Sharma

DIFFICULTY:

3199

PREREQUISITES:

Breadth First Search, Depth First Search, Probability

PROBLEM:

Alice and Bob live in a country in which there are N cities (numbered from 1 to N) connected to each other using N-1 bidirectional roads in such a way that any two cities are reachable from each other.

This country has a special property: for each city in the country, there are at most 6 roads directly connected to it.

Alice and Bob both decide to go on a tour of the country, starting from city number 1. Alice uses Depth First traversal (DFS) order to travel through all the N cities while Bob uses Breadth First Traversal (BFS) order. At any city, if there is more than one possibility for the next city to be visited, then each option is equally likely to be chosen by both Alice and Bob.

The pseudocode of their respective traversals are as follows:

// Alice's DFS traversal
dfs(u):
    // Let C denote the set of children of u
    shuffle(C)
    // Let the current state of C be [C[1], C[2], ..., C[k]]
    for i from 1 to k:
        dfs(C[i])
// Bob's BFS traversal
queue Q
Q.push(1)
while Q is not empty:
    u = front(Q)
    pop(Q)
    // Let C denote the set of children of u
    shuffle(C)
    // Let the current state of C be [C[1], C[2], ..., C[k]]
    for i from 1 to k:
        Q.push(C[i])

Define the index of a city in some traversal order to be the number of distinct cities visited before that city during that traversal. Find the expected number of cities whose index in Bob’s traversal is strictly less than its index in Alice’s traversal.

It can be shown that this expected value can be expressed as a fraction \frac{P}{Q}, where P and Q are coprime integers, P \ge 0, Q \gt 0 and Q is coprime with 998244353. You should compute (P \cdot Q^{-1}) \pmod {998\,244\,353}, where Q^{-1} denotes the multiplicative inverse of Q modulo 998\,244\,353.

EXPLANATION:

Let D[i][j] be the probability that i_{th} city is at the j_{th} index in the DFS traversal, and B[i][j] represent the same for BFS traversal. First of all, we will compute the matrices D and B. We can assume the given graph network as a tree rooted at node 1.

Computation of D :

Suppose we know D[p][0], D[p][1], ..., D[p][N-1] for some node p. Also, let (c_1, c_2,..,c_m) be the children of p. Now using the probabilities for p, we will compute the probabilities for each of its child.
There will be in total m! ways in which Alice can travel the children of p. Therefore, probability of occurrence of each permutation is \frac{1}{m!}. Without loss of generality, let us assume a particular order, say (c_1,c_2,..,c_m). If the index of p is suppose k in the DFS order then, the index of child c_i , id[c_i] will be -

id[c_i] = k+1 + \sum_{j=1}^{i-1}SubtreeSize[c_j], where SubtreeSize[j] is the size of subtree rooted at j.

So, we can update D as-

D[c_i][id[c_i]] += \frac{1}{m!}*D[p][k], for (1 \le i \le m) and (0 \le k \lt N)

Node 1 will always have index 0, so using this base case, we can compute the complete D matrix.

Computation of B :

Let p be a node in the tree and (c_1,c_2,..,c_m) be its children. Suppose, we have computed B matrices for the subtree with each of it’s children as root separately. Let them be named as B_{c_1}, B_{c_2},..B_{c_m}. Now, we will compute B matrix for the subtree rooted at p, i.e. B_p.
There will be in total m! ways in which Bob can travel the children of p. Therefore, probability of occurrence of each permutation is \frac{1}{m!}.
Without loss of generality, Let us assume a particular order, say (c_1,c_2,..,c_m). Now consider any node x at some level y in the subtree of node c_i. In the BFS traversal of the subtree rooted at p, index of x will always lie in the range L to R, where

L = \sum_{j=0}^{j=y}level\_cnt[p][j] + \sum_{j=1}^{j=i-1}level\_cnt[c_j][y],
R = L + level\_cnt[c_i][y]

Here, level\_cnt[a][b] denotes the count of nodes at a distance of b (or at level b) in the subtree of node a.
Also, in the BFS traversal of the subtree rooted at node c_i, node x will always lie in the range L' and R', where L' and R' are computed in similar way as L and R respectively for the subtree of c_i instead that of p. So, we can compute B_p in the folowing manner-

B_p[x][L+j] += \frac{1}{m!}*B_{c_i}[x][L'+j] ~~\forall~ j, 0 \le j \le R-L

We need to update for all the nodes x in the subtree of node p for all possible m! permutation of it’s children.

Here the base case is when we are at the leaf node, in such case B_p[p][0] = 1 and B_p[p][i] = 0 for all i>0. And the final B matrix is same as B_1.

After Computation of B and D matrices, we can compute the required expected value (E) as follows-

E = \sum_{i=1}^{i=N} \text{(Probability that index of node i in BFS order is less than that in DFS order)}

E = \sum_{i=1}^{i=N} \left[\sum_{j=0}^{j=N-1}\left(B[i][j]*\sum_{k=j+1}^{k=N-1}D[i][k]\right)\right]

TIME COMPLEXITY:

O(N^3) per test case

SOLUTION:

Setter's/Editorialist's solution
#include <bits/stdc++.h>
using namespace std;
 
const int MAX_T = 1e5;
const int MAX_N = 1e5;
const int MAX_SUM_LEN = 1e5;
 
#define fast ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0)
#define ff first
#define ss second
#define mp make_pair
#define ll long long
#define rep(i,n) for(int i=0;i<n;i++)
#define rev(i,n) for(int i=n;i>=0;i--)
#define rep_a(i,a,n) for(int i=a;i<n;i++)
#define pb push_back
#define all(ds) ds.begin(), ds.end()
#define int ll


ll mod = 998244353;

ll po(ll x, ll n ){ 
    ll ans=1;
     while(n>0){
        if(n&1) ans=(ans*x)%mod;
        x=(x*x)%mod;
        n/=2;
     }
     return ans;
}

vector<ll> fac_inv(10);

void pre(){
    fac_inv[0]=1;
    rep_a(i,1,10) fac_inv[i]=(i*fac_inv[i-1]);
    rep(i,10){
        fac_inv[i] = po(fac_inv[i], mod-2);
    }
}

vector<vector<int> > adj, lvl_cnt;
vector<int> subsz;

vector<vector<ll> > dfsp, bfsp, comp;
int cnt;
int n;

void chk(int c, int p){
    cnt++;
    subsz[c]=1;
    lvl_cnt[c][0]=1;
    for(auto h:adj[c]){
        if(h!=p){
            chk(h,c);
            subsz[c]+=subsz[h];
            rep_a(i,1,n) lvl_cnt[c][i]+=lvl_cnt[h][i-1];
        }
    }
}

///////////////compute DFS-array///////////////


void dfsC(int c, int p){
    vector<int> tmp;
    for(auto h:adj[c]){
        if(h!=p) tmp.pb(h);
    }

    sort(all(tmp));
    int z = tmp.size();
    if(z==0) return;

    do{
        int curr = 1;
        for(auto h:tmp){
            for(int i=0; i+curr<n; i++){
                dfsp[h][i+curr]+=dfsp[c][i];
            }

            curr+=subsz[h];
        }

    }while(next_permutation(all(tmp)));

    for(auto h:adj[c]){
        if(h!=p){
            rep(i,n){
                dfsp[h][i] = ((dfsp[h][i]%mod)*fac_inv[z])%mod;
            }

            dfsC(h,c);
        }
    }


}


/////////////////////////////////////////////

///////////////compute BFS-array///////////////

void fun(vector<int>&y, vector<int>&y1, int c, int p, int lvl){  ///update subtree position


    rep_a(i,y[lvl], n){
        comp[c][i]+=bfsp[c][i-y[lvl]+(lvl>0?y1[lvl-1]:0)];
    }
    for(auto h:adj[c]){
        if(h!=p) fun(y,y1,h,c,lvl+1);
    }
}

void fun1(int c, int p, ll den){  /// copy from temporary array and turn into probability
    bfsp[c] = comp[c];
    comp[c].assign(n,0);
    rep(i,n){
        bfsp[c][i] = ((bfsp[c][i]%mod)*den)%mod;
    }
    for(auto h:adj[c]){
        if(h!=p) fun1(h,c,den);
    }
}


void bfsC(int c, int p){  //////consider each permutation for all children 

    vector<int> tmp;

    for(auto h:adj[c]){
        if(h!=p){
            bfsC(h,c);
            tmp.pb(h);
        }
    }

    sort(all(tmp));

    
    bfsp[c][0]=1;

    if(tmp.empty()) return;
    ll z = tmp.size();

    do{
        vector<int> y = lvl_cnt[c];
        rep_a(i,1,n) y[i]+=y[i-1];
        vector<int> y1;

        for(auto h:tmp){
            y1 = lvl_cnt[h];
            rep_a(i,1,n) y1[i]+=y1[i-1];
            fun(y,y1, h,c,0);

            rep(i,n) y[i]+=lvl_cnt[h][i];
        }
    }while(next_permutation(all(tmp)));

    for(auto h:adj[c]){
        if(h!=p){
            fun1(h,c,fac_inv[z]);
        }
    }


       


    
}

///////////////////////////////////////////////////////

void solve(){
    cin >> n;

    adj.assign(n, vector<int>());
    lvl_cnt.assign(n, vector<int>(n,0));
    comp.assign(n, vector<ll>(n,0));
    subsz.assign(n,0);

    dfsp.assign(n, vector<ll>(n,0));
    bfsp.assign(n, vector<ll>(n,0));

    int x,y;

    rep(i,n-1){
        cin>>x>>y;
        x--, y--;

        adj[x].pb(y);
        adj[y].pb(x);
    }

    cnt=0;
    chk(0,-1);
    assert(cnt==n);

    dfsp[0][0]=1;
    dfsC(0,-1);

    bfsC(0,-1);

    ll ans = 0;
    rep(i,n){
        ll dz = 0;
        rep(j,n) dz+=dfsp[i][j];
        dz%=mod;
        rep(j,n){
            dz-=dfsp[i][j];
            dz%=mod;
            ans += (bfsp[i][j]*dz)%mod;
        }
        ans%=mod;
    }

    if(ans<0) ans+=mod;
    cout<<ans<<'\n';


}
 
signed main()
{

    #ifndef ONLINE_JUDGE
    freopen("input.txt", "r" , stdin);
    freopen("output.txt", "w" , stdout);
    #endif
    fast;
    
    int t = 1;
    pre();

   cin>>t;
    
    for(int i=1;i<=t;i++)
    {    
       solve();
    }

    cerr << "Time : " << 1000 * ((double)clock()) / (double)CLOCKS_PER_SEC << "ms\n";
}

Tester's solution
#include "bits/stdc++.h"
// #pragma GCC optimize("O3,unroll-loops")
// #pragma GCC target("avx2,bmi,bmi2,lzcnt,popcnt")
using namespace std;
using ll = long long int;
mt19937_64 rng(chrono::high_resolution_clock::now().time_since_epoch().count());

/**
 * Integers modulo p, where p is a prime
 * Source: Aeren (modified from tourist?)
 *         Modmul for 64-bit mod from kactl:ModMulLL
 * Works with p < 7.2e18 with x87 80-bit long double, and p < 2^52 ~ 4.5e12 with 64-bit
 */
template<typename T>
struct Z_p{
	using Type = typename decay<decltype(T::value)>::type;
	static vector<Type> MOD_INV;
	constexpr Z_p(): value(){ }
	template<typename U> Z_p(const U &x){ value = normalize(x); }
	template<typename U> static Type normalize(const U &x){
		Type v;
		if(-mod() <= x && x < mod()) v = static_cast<Type>(x);
		else v = static_cast<Type>(x % mod());
		if(v < 0) v += mod();
		return v;
	}
	const Type& operator()() const{ return value; }
	template<typename U> explicit operator U() const{ return static_cast<U>(value); }
	constexpr static Type mod(){ return T::value; }
	Z_p &operator+=(const Z_p &otr){ if((value += otr.value) >= mod()) value -= mod(); return *this; }
	Z_p &operator-=(const Z_p &otr){ if((value -= otr.value) < 0) value += mod(); return *this; }
	template<typename U> Z_p &operator+=(const U &otr){ return *this += Z_p(otr); }
	template<typename U> Z_p &operator-=(const U &otr){ return *this -= Z_p(otr); }
	Z_p &operator++(){ return *this += 1; }
	Z_p &operator--(){ return *this -= 1; }
	Z_p operator++(int){ Z_p result(*this); *this += 1; return result; }
	Z_p operator--(int){ Z_p result(*this); *this -= 1; return result; }
	Z_p operator-() const{ return Z_p(-value); }
	template<typename U = T>
	typename enable_if<is_same<typename Z_p<U>::Type, int>::value, Z_p>::type &operator*=(const Z_p& rhs){
		#ifdef _WIN32
		uint64_t x = static_cast<int64_t>(value) * static_cast<int64_t>(rhs.value);
		uint32_t xh = static_cast<uint32_t>(x >> 32), xl = static_cast<uint32_t>(x), d, m;
		asm(
			"divl %4; \n\t"
			: "=a" (d), "=d" (m)
			: "d" (xh), "a" (xl), "r" (mod())
		);
		value = m;
		#else
		value = normalize(static_cast<int64_t>(value) * static_cast<int64_t>(rhs.value));
		#endif
		return *this;
	}
	template<typename U = T>
	typename enable_if<is_same<typename Z_p<U>::Type, int64_t>::value, Z_p>::type &operator*=(const Z_p &rhs){
		uint64_t ret = static_cast<uint64_t>(value) * static_cast<uint64_t>(rhs.value) - static_cast<uint64_t>(mod()) * static_cast<uint64_t>(1.L / static_cast<uint64_t>(mod()) * static_cast<uint64_t>(value) * static_cast<uint64_t>(rhs.value));
		value = normalize(static_cast<int64_t>(ret + static_cast<uint64_t>(mod()) * (ret < 0) - static_cast<uint64_t>(mod()) * (ret >= static_cast<uint64_t>(mod()))));
		return *this;
	}
	template<typename U = T>
	typename enable_if<!is_integral<typename Z_p<U>::Type>::value, Z_p>::type &operator*=(const Z_p &rhs){
		value = normalize(value * rhs.value);
		return *this;
	}
	template<typename U>
	Z_p &operator^=(U e){
		if(e < 0) *this = 1 / *this, e = -e;
		Z_p res = 1;
		for(; e; *this *= *this, e >>= 1) if(e & 1) res *= *this;
		return *this = res;
	}
	template<typename U>
	Z_p operator^(U e) const{
		return Z_p(*this) ^= e;
	}
	Z_p &operator/=(const Z_p &otr){
		Type a = otr.value, m = mod(), u = 0, v = 1;
		if(a < (int)MOD_INV.size()) return *this *= MOD_INV[a];
		while(a){
			Type t = m / a;
			m -= t * a; swap(a, m);
			u -= t * v; swap(u, v);
		}
		assert(m == 1);
		return *this *= u;
	}
	template<typename U> friend const Z_p<U> &abs(const Z_p<U> &v){ return v; }
	Type value;
};
template<typename T> bool operator==(const Z_p<T> &lhs, const Z_p<T> &rhs){ return lhs.value == rhs.value; }
template<typename T, typename U, typename enable_if<is_integral<U>::value>::type* = nullptr> bool operator==(const Z_p<T>& lhs, U rhs){ return lhs == Z_p<T>(rhs); }
template<typename T, typename U, typename enable_if<is_integral<U>::value>::type* = nullptr> bool operator==(U lhs, const Z_p<T> &rhs){ return Z_p<T>(lhs) == rhs; }
template<typename T> bool operator!=(const Z_p<T> &lhs, const Z_p<T> &rhs){ return !(lhs == rhs); }
template<typename T, typename U, typename enable_if<is_integral<U>::value>::type* = nullptr> bool operator!=(const Z_p<T> &lhs, U rhs){ return !(lhs == rhs); }
template<typename T, typename U, typename enable_if<is_integral<U>::value>::type* = nullptr> bool operator!=(U lhs, const Z_p<T> &rhs){ return !(lhs == rhs); }
template<typename T> bool operator<(const Z_p<T> &lhs, const Z_p<T> &rhs){ return lhs.value < rhs.value; }
template<typename T> bool operator>(const Z_p<T> &lhs, const Z_p<T> &rhs){ return lhs.value > rhs.value; }
template<typename T> bool operator<=(const Z_p<T> &lhs, const Z_p<T> &rhs){ return lhs.value <= rhs.value; }
template<typename T> bool operator>=(const Z_p<T> &lhs, const Z_p<T> &rhs){ return lhs.value >= rhs.value; }
template<typename T> Z_p<T> operator+(const Z_p<T> &lhs, const Z_p<T> &rhs){ return Z_p<T>(lhs) += rhs; }
template<typename T, typename U, typename enable_if<is_integral<U>::value>::type* = nullptr> Z_p<T> operator+(const Z_p<T> &lhs, U rhs){ return Z_p<T>(lhs) += rhs; }
template<typename T, typename U, typename enable_if<is_integral<U>::value>::type* = nullptr> Z_p<T> operator+(U lhs, const Z_p<T> &rhs){ return Z_p<T>(lhs) += rhs; }
template<typename T> Z_p<T> operator-(const Z_p<T> &lhs, const Z_p<T> &rhs){ return Z_p<T>(lhs) -= rhs; }
template<typename T, typename U, typename enable_if<is_integral<U>::value>::type* = nullptr> Z_p<T> operator-(const Z_p<T>& lhs, U rhs){ return Z_p<T>(lhs) -= rhs; }
template<typename T, typename U, typename enable_if<is_integral<U>::value>::type* = nullptr> Z_p<T> operator-(U lhs, const Z_p<T> &rhs){ return Z_p<T>(lhs) -= rhs; }
template<typename T> Z_p<T> operator*(const Z_p<T> &lhs, const Z_p<T> &rhs){ return Z_p<T>(lhs) *= rhs; }
template<typename T, typename U, typename enable_if<is_integral<U>::value>::type* = nullptr> Z_p<T> operator*(const Z_p<T>& lhs, U rhs){ return Z_p<T>(lhs) *= rhs; }
template<typename T, typename U, typename enable_if<is_integral<U>::value>::type* = nullptr> Z_p<T> operator*(U lhs, const Z_p<T> &rhs){ return Z_p<T>(lhs) *= rhs; }
template<typename T> Z_p<T> operator/(const Z_p<T> &lhs, const Z_p<T> &rhs) { return Z_p<T>(lhs) /= rhs; }
template<typename T, typename U, typename enable_if<is_integral<U>::value>::type* = nullptr> Z_p<T> operator/(const Z_p<T>& lhs, U rhs) { return Z_p<T>(lhs) /= rhs; }
template<typename T, typename U, typename enable_if<is_integral<U>::value>::type* = nullptr> Z_p<T> operator/(U lhs, const Z_p<T> &rhs) { return Z_p<T>(lhs) /= rhs; }
template<typename T> istream &operator>>(istream &in, Z_p<T> &number){
	typename common_type<typename Z_p<T>::Type, int64_t>::type x;
	in >> x;
	number.value = Z_p<T>::normalize(x);
	return in;
}
template<typename T> ostream &operator<<(ostream &out, const Z_p<T> &number){ return out << number(); }

/*
using ModType = int;
struct VarMod{ static ModType value; };
ModType VarMod::value;
ModType &mod = VarMod::value;
using Zp = Z_p<VarMod>;
*/

// constexpr int mod = 1e9 + 7; // 1000000007
constexpr int mod = (119 << 23) + 1; // 998244353
// constexpr int mod = 1e9 + 9; // 1000000009
using Zp = Z_p<integral_constant<decay<decltype(mod)>::type, mod>>;

template<typename T> vector<typename Z_p<T>::Type> Z_p<T>::MOD_INV;
template<typename T = integral_constant<decay<decltype(mod)>::type, mod>>
void precalc_inverse(int SZ){
	auto &inv = Z_p<T>::MOD_INV;
	if(inv.empty()) inv.assign(2, 1);
	for(; inv.size() <= SZ; ) inv.push_back((mod - 1LL * mod / (int)inv.size() * inv[mod % (int)inv.size()]) % mod);
}

template<typename T>
vector<T> precalc_power(T base, int SZ){
	vector<T> res(SZ + 1, 1);
	for(auto i = 1; i <= SZ; ++ i) res[i] = res[i - 1] * base;
	return res;
}

template<typename T>
vector<T> precalc_factorial(int SZ){
	vector<T> res(SZ + 1, 1); res[0] = 1;
	for(auto i = 1; i <= SZ; ++ i) res[i] = res[i - 1] * i;
	return res;
}

int main()
{
	ios::sync_with_stdio(false); cin.tie(0);

	auto fac = precalc_factorial<Zp>(10);

	int t; cin >> t;
	while (t--) {
		int n; cin >> n;
		vector<vector<int>> adj(n);
		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<int> subsz(n), childct(n);
		auto calc_sz = [&] (const auto &self, int u, int par) -> void {
			subsz[u] = 1;
			for (int v : adj[u]) {
				if (v == par) continue;
				self(self, v, u);
				subsz[u] += subsz[v];
				++childct[u];
			}
		};
		
		vector dp_dfs(n, vector(n+5, Zp(0)));
		auto calc_dfs = [&] (const auto &self, int u, int par) -> void {
			vector<int> children;
			for (int v : adj[u]) {
				if (v == par) continue;
				children.push_back(v);
			}
			const Zp prob = Zp(1)/fac[childct[u]];
			for (int i = 0; i < n; ++i) {
				if (dp_dfs[u][i] == 0) continue;
				sort(begin(children), end(children));
				do {
					int add = 0;
					for (int v : children) {
						dp_dfs[v][i+add+1] += prob*dp_dfs[u][i];
						add += subsz[v];
					}
				} while(next_permutation(begin(children), end(children)));
			}
			for (int v : adj[u]) {
				if (v == par) continue;
				self(self, v, u);
			}
		};
		
		vector dp_bfs(n, vector(n+5, Zp(0)));
		vector dp_bfs_new(n, vector(n+5, Zp(0)));
		auto calc_bfs = [&] (const auto &self, int u, int par) -> void {
			vector<int> children;
			for (int v : adj[u]) {
				if (v == par) continue;
				children.push_back(v);
				self(self, v, u);
			}
			const Zp prob = Zp(1)/fac[childct[u]];
			sort(begin(children), end(children));
			vector<int> subtree;
			do {
				queue<array<int, 4>> q;
				for (int v : children) q.push({v, v, 1, u});
				int before = 0, prv = -1, prvdep = 0;
				vector<int> active;
				map<int, int> en;
				while (!q.empty()) {
					auto [v, root, dep, p] = q.front(); q.pop();
					if (dep == prvdep and root == prv) active.push_back(v);
					else {
						int lo = before + 1, hi = before + active.size();
						int st = en[prv] - active.size();
						for (int i = lo; i <= hi; ++i) {
							for (int x : active) {
								dp_bfs_new[x][i] += prob*dp_bfs[x][st];
							}
							++st;
						}
						before += active.size();
						active = {v};
						prv = root;
						prvdep = dep;
					}
					for (int x : adj[v]) {
						if (x == p) continue;
						q.push({x, root, dep+1, v});
					}
					++en[root];

					if (is_sorted(begin(children), end(children))) subtree.push_back(v);
				}
				int lo = before + 1, hi = before + active.size();
				int st = en[prv] - active.size();
				for (int i = lo; i <= hi; ++i) {
					for (int x : active) {
						dp_bfs_new[x][i] += prob*dp_bfs[x][st];
					}
					++st;
				}
			} while(next_permutation(begin(children), end(children)));
			
			dp_bfs[u][0] = 1;
			for (int v : subtree) {
				swap(dp_bfs[v], dp_bfs_new[v]);
				dp_bfs_new[v].assign(n+5, 0);
			}

		};
		calc_sz(calc_sz, 0, 0);
		dp_dfs[0][0] = 1;
		calc_dfs(calc_dfs, 0, 0);
		calc_bfs(calc_bfs, 0, 0);
		Zp ans = 0;
		for (int i = 0; i < n; ++i) {
			Zp pref = 0;
			for (int j = 0; j < n; ++j) {
				// cerr << dp_bfs[i][j] << ' ';
				ans += dp_dfs[i][j] * pref;
				pref += dp_bfs[i][j];
			}
			// cerr << endl;
		}
		cout << ans << '\n';
	}
}