PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: tanminati
Tester: sushil2006
Editorialist: iceknight1093
DIFFICULTY:
Easy
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 twice.
EXPLANATION:
There’s a unique simple path in the tree from 1 to N.
Let this be v_1 \to v_2 \to\ldots\to v_k, where v_1 = 1 and v_k = N.
We start at v_1 = 1.
Suppose the first move of the walk is to move to vertex x.
There are two possibilities: either x = v_2 or x \ne v_2.
First, let’s look at x = v_2.
Once we move to v_2, there are a further two possibilities: either we turn back to v_1 immediately, or we don’t.
Note that if we don’t immediately turn back to v_1, then we can also never touch v_1 in the future either - since we entered v_2 once, then to move to v_1 we’ll need to re-enter v_2, and then from v_1 we’ll eventually need to move to v_2 again to reach N; leading to three visits to v_2 which is not allowed.
Similar reasoning shows that if we immediately move to v_1, then the next move must be to v_2, and then next immediately to v_3; after which we’re basically starting afresh at v_3 since we can no longer move to v_2.
So, suppose we define a function f(v_i) denoting the number of walks reaching N that start from v_i, assuming that we’re at v_i for the first time and will never go back to v_{i-1} during the walk.
Our aim is to compute f(1).
Then, if the first move is v_1 \to v_2, the above analysis tells us that:
- There are f(v_2) walks which don’t immediately go to v_1 after going to v_2.
- There are f(v_3) walks which immediately go to v_1 after going to v_2; since they’re forced to start as v_1 \to v_2 \to v_1 \to v_2 \to v_3.
Next, let’s analyze the case where the first move is not v_1 \to v_2, i.e. we move to some child of v_1 that’s not v_2.
Suppose we move to a child of v_1.
From this child, we can again continue moving ‘downwards’ through the tree; till eventually we make our first upward move (which must happen, since we need to eventually return to 1).
Observe that after the first upward move, all further moves are forced to be upward too till we reach vertex 1 again.
This is because every vertex visited during an upward move will be visited a second time; and making a downward move from it will then require visiting it a third time to return to 1 which is not allowed.
After reaching v_1 = 1, our next move is then forced to be to v_2.
Thus, the walk will look like: choose some vertex x such that the path 1 \to x doesn’t include v_2.
Then, move along the path from 1 to x, then move back to 1 along this path, then move to v_2 and continue from there.
Each valid vertex x gives us exactly one valid beginning to the walk.
So, if there are c valid vertices, the number of walks obtained in this case is f(v_2) \cdot c.
How to compute c?
Well, it’s quite simply all those vertices that are not present in the subtree of v_2 (and also not including 1).
This equals N - 1 - \text{subsz}[v_2], where \text{subsz}[v_2] is the size of the subtree rooted at v_2.
Putting both cases together, we have:
It’s not hard to see that this applies to a general v_i as well.
The only difference is the N - 1 - \text{subsz}[v_2] term, which becomes \text{subsz}[v_{i}] - 1 - \text{subsz}[v_{i+1}] instead, so we obtain
You might need to handle the case of i+1 = k specially, since it’s not allowed to move back from v_k = N after reaching it; this only changes the transitions slightly for this one case.
After computing subtree sizes and the 1\to N, we can thus compute the necessary answer in linear time using dynamic programming.
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;
// for each node on path, find #of nodes in subtree (note: can never step on node that is on the path)
vector<ll> subcnt(n+5);
auto dfs2 = [&](ll u, ll p, auto &&dfs2) -> void{
subcnt[u] = 1;
trav(v,adj[u]){
if(v == p or on_path[v]) conts;
dfs2(v,u,dfs2);
subcnt[u] += subcnt[v];
}
};
trav(u,path){
dfs2(u,-1,dfs2);
}
// do the dp
ll siz = sz(path);
vector<ll> dp(siz+5);
dp[0] = 1; // start from node 1 (index 0 on the path)
rep(i,siz){
ll u = path[i];
dp[i] %= MOD;
// 1) go down some subtree of u (could also stay at u only, is included in this case), then come back upto u, then go to (i+1)th node on the path --> ways = subcnt[u]
dp[i+1] += dp[i]*subcnt[u];
// 2) u --> next(u) --> u --> next(u) --> next(next(u)) i.e i --> i+2 with exactly 1 way
dp[i+2] += dp[i];
}
ll ans = dp[siz-1]%MOD;
cout << ans << endl;
}
int main()
{
fastio;
int t = 1;
cin >> t;
rep1(i, t) {
solve(i);
}
cerr << "RUN SUCCESSFUL" << endl;
return 0;
}