PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: tanminati
Tester: sushil2006
Editorialist: iceknight1093
DIFFICULTY:
Medium
PREREQUISITES:
DFS, dynamic programming
PROBLEM:
You’re given a tree.
Count the number of walks that start at 1, and at N, visit N exactly once, and visit any other vertex at most thrice.
EXPLANATION:
Just as with the easy version, we need to analyze the possible walks.
However, while in the easy version walks were highly constrained, being able to visit vertices three times gives us much more freedom here.
Let v_1 \to v_2\to \ldots \to v_k be the path from 1 to N, where v_1 = 1 and v_k = N.
We must move along this path; but we can also move back-and-forth along it a bit, and step outside of it here and there.
Note that for each ‘outside’ node u, if we visit it and go into its subtree, we also need to return to it eventually to be able to return to the main path.
Further, if we go to a node, visit its subtree, then leave the node (as in move to its parent) and then return to it later, we cannot go into its subtree again (since to leave it again will require more than three visits.)
Thus, we can go into a subtree non-trivially at most once.
So, let’s define out(u, x) to be the number of walks that start at u (given that this is the first time visiting u), go into the subtree of u and eventually end up at u with the final number of visits to u being x.
(Note that we allow for u being visited in the future from its parent - that’s not being counted here, and will be handled from its parent instead.)
Computing the out(u, x) values requires some combinatorics.
- out(u, 1) is trivially 1, since we can’t actually go into the subtree of u and must just remain here.
- For out(u, 2), we must move to one child subtree of u, explore it in some way, and then return to u.
The number of ways here equals the sum of all out(c, y) across all children c of u and choices of y. - out(u, 3) requires the most effort to compute, since there are multiple cases.
The following walks are possible:- Move into one child subtree, explore it, return to u.
Then move into a different subtree, explore it, return to u. - Move into one child subtree, explore it, return to u.
Then move to the same child and return immediately to u (which requires the child to be visited \lt 3 times.)
Alternately, reverse the order: move to the child, return, then move to the same child and explore it before returning.
- Move into one child subtree, explore it, return to u.
out(u, 3) can also be computed from the out(c, y) values of the children c of u.
The exact transitions are left as an exercise to the reader - or see the attached code.
In any case, we can compute out(u, x) for all u and x (where u does not lie on the 1\to N path) in linear time.
With the helper function above, we can now move on to a full solution for the problem.
Let’s look at the main 1 \to N path, and specifically how we move along this path.
For now, we ignore the parts of the walk that are outside this main path (since each such outside part will return to the main path eventually, just adding 1 to the number of occurrences of a path vertex.)
We can kind of move back-and-forth along the main path.
That is, we start at v_1, and initially our moves will look like v_1 \to v_2 \to v_3 \to\ldots
However, eventually we’ll make a backward move, i.e. v_i \to v_{i-1} for some i.
Once this happens, there are two options for what we can do:
- Continue moving backwards a bit till we reach some v_j where j \lt i.
Then, move forward once again.
Observe that in this case, as soon as we start moving forward, we’re forced to move all the way forward till we reach v_{i+1}, since all the vertices with indices in [j+1, i-1] will have three visits to them.
Note that after this happens, a future backward movement can only reach upto v_i at most, and cannot go further than that. - Alternately, we can also ‘oscillate’ between v_i and v_{i-1}, i.e v_{i-1} \to v_i \to v_{i-1} \to v_i \to v_{i-1} \to v_i, after which we’re forced to move to v_{i+1} again.
In this case, future backward moves cannot even reach v_i, and must stop at v_{i+1}.
With this in mind, let’s define dp(i, j) to be the number of walks that are currently at v_i, v_i has been visited j times, and v_{i+1} has not been visited yet.
Further, we’re still only considering walks along the main path for now - we’ll bring in outside stuff later.
The transitions are not too hard.
- First, dp(i, 1).
This means we have to come here from v_{i-1}, and it doesn’t matter how many times that one was visited.
So, just sum up all dp(i-1, j) to obtain this. - Next, dp(i, 2).
For this, note that we must first visit v_i, then double back till we reach some previous v_j, and then return to i.
Further, all vertices with indices in [j+1, i-1] must’ve been visited only once along this walk; and v_j must’ve been visited either once or twice.
So, for each j \lt i, there are dp(j, 1) + dp(j, 2) possible ways to obtain such a walk.
This can be computed quickly using prefix sums. - Finally, dp(i, 3).
The only way to obtain this is to reach v_{i-1} with a count of 1, then oscillate between v_{i-1} and v_i.
So, the number of ways here is dp(i-1, 1)
Note that we don’t even need to consider this state - since the next move forces us to move to v_{i+1} anyway, we can just club this transition into dp(i+1, 1) instead (which is what the attached code does.)
Given that we ignored all outside parts, this is a linear-time DP.
Now, let’s try to incorporate outside parts as well.
Observe that these outside parts can be included at any point when we visit a vertex.
However, their number is limited: if a path vertex is visited three times along the walk already, no outside part involving it can be done.
If the vertex is visited twice, (up to) one outside part is possible; and can be included either after the first visit or after the second.
If the vertex is visited only once, (up to) two outside parts are possible, and both must be done immediately upon visiting it.
With these observations (and a little combinatorics for how to interleave outer parts), and careful observation of which vertices along the path have which visited counts when looking at each of the transitions, it’s possible to incorporate everything the outer paths into our above dp(i, j) definition via appropriate multipliers, while keeping the complexity linear.
For precise formulas, please view the attached code.
TIME COMPLEXITY:
\mathcal{O}(N) per testcase.
CODE:
Tester's code (C++)
#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace std;
using namespace __gnu_pbds;
template<typename T> using Tree = tree<T, null_type, less<T>, rb_tree_tag, tree_order_statistics_node_update>;
typedef long long int ll;
typedef long double ld;
typedef pair<int,int> pii;
typedef pair<ll,ll> pll;
#define fastio ios_base::sync_with_stdio(false); cin.tie(NULL)
#define pb push_back
#define endl '\n'
#define sz(a) (int)a.size()
#define setbits(x) __builtin_popcountll(x)
#define ff first
#define ss second
#define conts continue
#define ceil2(x,y) ((x+y-1)/(y))
#define all(a) a.begin(), a.end()
#define rall(a) a.rbegin(), a.rend()
#define yes cout << "Yes" << endl
#define no cout << "No" << endl
#define rep(i,n) for(int i = 0; i < n; ++i)
#define rep1(i,n) for(int i = 1; i <= n; ++i)
#define rev(i,s,e) for(int i = s; i >= e; --i)
#define trav(i,a) for(auto &i : a)
template<typename T>
void amin(T &a, T b) {
a = min(a,b);
}
template<typename T>
void amax(T &a, T b) {
a = max(a,b);
}
#ifdef LOCAL
#include "debug.h"
#else
#define debug(...) 42
#endif
/*
*/
const int MOD = 998244353;
const int N = 1e5 + 5;
const int inf1 = int(1e9) + 5;
const ll inf2 = ll(1e18) + 5;
void solve(int test_case){
ll n; cin >> n;
vector<vector<ll>> adj(n+5);
rep1(i,n-1){
ll u,v; cin >> u >> v;
adj[u].pb(v), adj[v].pb(u);
}
// extract the 1 <--> n path
vector<ll> curr_path, path;
auto dfs1 = [&](ll u, ll p, auto &&dfs1) -> void{
curr_path.pb(u);
if(u == n){
path = curr_path;
}
trav(v,adj[u]){
if(v == p) conts;
dfs1(v,u,dfs1);
}
curr_path.pop_back();
};
dfs1(1,-1,dfs1);
// mark all nodes on the path
vector<bool> on_path(n+5);
trav(u,path) on_path[u] = 1;
// ways[u][j] = #of ways to finish the sub of u if u will be visited "j MORE" times after this
// note: never step on node that is on the path
vector<vector<ll>> ways(n+5,vector<ll>(5));
vector<ll> ways_sum(n+5);
auto dfs2 = [&](ll u, ll p, auto &&dfs2) -> void{
vector<ll> children;
trav(v,adj[u]){
if(v == p or on_path[v]) conts;
dfs2(v,u,dfs2);
children.pb(v);
}
// ways[u][0]
// do nothing --> 1
ways[u][0] = 1;
// ways[u][1]
// go down to 1 subtree, return back to u --> ways_sum[v]
trav(v,children){
ways[u][1] += ways_sum[v];
}
ways[u][1] %= MOD;
// ways[u][2]
// go to v1, back to u, go to v2, back to u (v1 != v2) --> ways_sum[v1]*ways_sum[v2]
// assume v1 < v2, then multiply tot sum by 2
ll sum = 0;
trav(v,children){
ways[u][2] += sum*ways_sum[v];
ways[u][2] %= MOD;
sum += ways_sum[v];
sum %= MOD;
}
ways[u][2] = ways[u][2]*2%MOD;
// both times go down into the same subtree
// one time just touch v and return, next time do anything inside v --> ways[v][0]+ways[v][1]
// do anything inside v, just touch v --> ways[v][0]+ways[v][1]
// overcounting: both times just touch v --> counted twice, but need it counted only once; subtract 1
trav(v,children){
ways[u][2] += (ways[v][0]+ways[v][1])*2-1+MOD;
ways[u][2] %= MOD;
}
// ways_sum[u]
ways_sum[u] = (ways[u][0]+ways[u][1]+ways[u][2])%MOD;
};
// compute ways dp for all nodes on the path
trav(u,path){
dfs2(u,-1,dfs2);
}
// do the main path dp
ll siz = sz(path);
ll dp[siz+5][5];
memset(dp,0,sizeof dp);
dp[0][1] = 1; // start at node 1 (idx 0 on the path), visited node 1 once so far
auto choose = [&](ll x, ll y) -> ll{
if(y == 0) return 1;
if(y == 1) return x;
if(y == 2) return x*(x-1)/2;
return -1;
};
auto fill_ways = [&](ll u, ll j){
// u has been visited j times so far (on the path)
// how many ways to fill the remaining
ll sum = 0;
rep(t,3){
// t = remaining vis count
if(j+t > 3) break;
ll w = ways[u][t];
ll c = choose(j+t-1,t);
sum += w*c;
}
return sum;
};
ll two_sum = 0;
rep(i,siz){
ll u = path[i];
rep1(j,2){
if(j == 2){
dp[i][j] = two_sum;
}
ll curr = dp[i][j];
if(!curr) conts;
// just go to next node on path
dp[i+1][1] += curr*fill_ways(u,j);
dp[i+1][1] %= MOD;
// go down to some v, come back up to u, go back down to the v
/*
for(int v = i+1; v < siz; ++v){
dp[v][2] += curr*fill_ways(u,j+1);
dp[v][2] %= MOD;
}
*/
// i --> i+1 --> i --> i+1 --> i --> i+1 --> i+2
if(j == 1){
dp[i+2][1] += curr;
dp[i+2][1] %= MOD;
}
}
rep1(j,2){
two_sum += dp[i][j]*fill_ways(u,j+1);
two_sum %= MOD;
}
}
ll ans = dp[siz-1][1];
cout << ans << endl;
}
int main()
{
fastio;
int t = 1;
cin >> t;
rep1(i, t) {
solve(i);
}
cerr << "RUN SUCCESSFUL" << endl;
return 0;
}