PROBLEM LINK:
Author: Arjun Bharat
Editorialist: Arjun Bharat
DIFFICULTY:
Hard
PREREQUISITES:
Graph algorithms, traversals, data structures, math
PROBLEM:
A tree is called 0-primal if it is an isolated root node.
A k-primal tree is one, which for a given fixed positive integer k, that consists of a root node that has a prime number of children,
each of which is the root of a k-1 primal tree.
Additionally, the number of prime children at each level of the tree must be unique, that is, for all non-leaf nodes at a given distance from the root,
the number of prime children they have must be unique.
Given a tree, check if it is k-primal or not, and if yes, identify the value of k that makes it k-primal.
SHORT EXPLANATION:
The key to the problem is identifying the root node of the tree, after which we can recursively check if the tree is k-primal. Naturally, we also have to identify the value of k. Once we have the root node and the value of k, a simple DFS will suffice to test for k-primality.
DETAILED EXPLANATION:
If we observe the property of the k-primal tree, we notice that it has a diameter of 2k + 1. Additionally the root node is the middle node of any diameter.
Convince yourself that these conditions are necessary and sufficient.
Using 2 BFSes, one from a random leaf node to find the farthest leaf,and one more from the farthest leaf, the diametric path can be found, and from here,
we know the value of k, as well as the root node.
To check for k-primality, we simply run a DFS from the root. We can recursively check for primality, while we ensure that each node has a prime number
of children. It would be useful to precompute and store primes using the sieve of Eratosthenes (upto 1.2*10^6).
Thus the overall complexity is O(1) + O(n) * 3 = O(n) (constant time to precompute primes, and 3 linear graph traversals)
ALTERNATIVE SOLUTION:
We also observe that for a k-primal tree, the root node must be at a distance k from any leaf node. Thus, if we run a BFS from all leaves simulataneously,
(by adding all of them into a queue initially and then doing a BFS), the first node whose every neighbour has been visited will be the root. Also noting down the distance of the root gives us the value of k. Again, the reader is requested to convince himself why this prpoerty is necessary and a sufficient check to discern the root.
The remaining part of the problem is the same as described earlier.
SOLUTION:
Setter's Solution
//bluishgreen
#pragma GCC optimize("-O3")
#include <bits/stdc++.h>
using namespace std;
#define REP(i,a,b) for(i=a;i!=b;i++)
#define getarr(a,n) for(i=0;i<n;i++){cin>>a[i];}
#define disparr(a,n) for(i=0;i<n;i++){cout<<a[i];}cout<<"\n";
#define disparrs(a,n) for(i=0;i<n;i++){cout<<a[i]<<" ";}cout<<"\n";
#define disparrl(a,n) for(i=0;i<n;i++){cout<<a[i]<<"\n";}cout<<"\n";
#define mem0(a) memset(a,0,sizeof(a));
#define ll long long int
#define ld long double
#define key pair<ld,ld>
#define pll pair<ll,ll>
#define pii pair<int,int>
#define si set<int>
#define vii vector<pair<int,int> >
#define vi vector<int>
#define vll vector<ll>
#define vb vector<bool>
#define vvi vector<vector<int>>
#define vs vector<string>
#define all(v) v.begin(),v.end()
#define pb push_back
#define pob pop_back
#define pof pop_front
#define mp make_pair
#define F first
#define S second
#define nu 100001
#define M 1000000007
#define mul(x,y) ((ll)(x)*(y))%M
#define tr(c,i) for(auto i = (c).begin(); i != (c).end(); i++)
#define fastio ios_base::sync_with_stdio(false); cin.tie(NULL); cout.tie(NULL)
#define INF 0x3f3f3f3f
#define flush fflush(stdout)
int MOD(int a, int b)
{
if(a>b)
return a-b;
else
return b-a;
}
ll max3(ll a,ll b,ll c)
{
return max(c,max(a,b));
}
ll power(ll x,ll y, ll p)
{
ll res = 1;
x = x % p;
while (y > 0)
{
if (y & 1)
res = (res*x) % p;
y = y>>1;
x = (x*x) % p;
}
return res;
}
ll max(ll a,ll b)
{
return (a>b)?a:b;
}
ll min(ll a,ll b)
{
return (a<b)?a:b;
}
ll gcd(ll a,ll b)
{
if(b==0)
return a;
else
return gcd(b,a%b);
}
ll gcd_ext(ll a,ll b,ll &x, ll &y)//solves a*x+b*y=gcd(a,b)
{
if(b==0)
{
x=1;
y=0;
return a;
}
else
{
ll tmp,x1,y1;
tmp = gcd_ext(b,a%b,x1,y1);//stores the result to b*x1+(a%b)*y1=gcd(b,a%b)=gcd(a,b)
x = y1;
y = x1 - (a/b)*y1;
return tmp;
}
}
vvi adj;
vvi children;
map<int,bool> prime;
vector<bool> vis;
map<int,map<int,bool>> occ;//this data structure does the uniqueness check
void build_tree(int s)
{
vis[s] = true;
for(auto it: adj[s])
{
if(!vis[it])
{
children[s].pb(it);
build_tree(it);
}
}
}
bool check(int s,int k)
{
if(k == 0)
return true;//0-primal case
int x = children[s].size();
map<int,bool> used;
if(!prime[x])
return false;//not primal if it doesn't have a prime number of children
if(occ[k][x])
return false;//the same prime repeats on a given level, hence not k-primal
occ[k][x] = true;
for(auto it: children[s])
{
if(check(it,k-1))
continue;
else
return false;//doesn't satisfy the recursive property.
}
return true;//k-primal
}
void eratosthenes()//to build a list of primes
{
int i,j;
vector<bool> visi(1200001,false);
for(i = 2;i <= 1200000;i++)
{
if(!visi[i])
{
prime[i] = true;
for(j = i;j <= 1200000;j += i)
{
visi[j] = true;
}
}
}
}
pii get_diam()
{
queue<int> q;
int i;
int n = vis.size();
for(i = 0;i < n;i++)
{
vis[i] = false;
if(adj[i].size() == 1 && q.empty())
q.push(i);
}
//cout<<"Hi!";
int src = q.front();
int level[n];
memset(level,-1,sizeof(level));
vis[src] = true;
level[src] = 0;
while(!q.empty())
{
int x = q.front();
q.pop();
for(auto it: adj[x])
{
if(!vis[it])
{
level[it] = level[x] + 1;
if(level[it] > level[src])
{
src = it;
}
vis[it] = true;
q.push(it);
}
}
}
//First BFS to find the farthest leaf is done
memset(level,-1,sizeof(level));
for(i = 0;i < n;i++)
vis[i] = false;
vis[src] = true;
int par[n];
memset(par,-1,sizeof(par));
q.push(src);
while(!q.empty())
{
int x = q.front();
q.pop();
for(auto it: adj[x])
{
if(!vis[it])
{
level[it] = level[x] + 1;
if(level[it] > level[src])
{
src = it;
}
par[it] = x;
vis[it] = true;
q.push(it);
}
}
}
//Second BFS to find the diameter is done
vector<int> v;
while(src != -1)
{
v.pb(src);
src = par[src];//find all nodes along the diameter
}
int k = (v.size()-1)/2;/*find the value of k: Note that even if the tree is not k-primal, we will get to know if a BFS from here*/
int centre = v[k];//central node
return mp(k,centre);
}
int main(void)
{
fastio;
eratosthenes();
int n;
cin>>n;
if(n == 1)
{
cout<<"YES\n0\n";
return 0;
}
adj.resize(n);
vis.resize(n);
children.resize(n);
int i,j,k;
for(i = 0;i < n-1;i++)
{
cin>>j>>k;
adj[j-1].pb(k-1);
adj[k-1].pb(j-1);
}
pii val = get_diam();
for(i=0;i<n;i++)
vis[i] = false;
build_tree(val.S);//identify children of each node by a DFS
if(val.F == 0)//if k is zero, the tree cannot be 0-primal as it has more than one node
{
cout<<"NO\n";
return 0;
}
if(check(val.S,val.F))
cout<<"YES\n"<<val.F<<'\n';
else
cout<<"NO\n";
return 0;
}