Root the tree at 1. We will do DP on subtrees.
First, observe that the answer is "no" if N−1 is not divisible by K. Otherwise, we wish to write a function dfs(x) which will check whether it is possible to partition the subtree corresponding to vertex x into paths of length K and possibly an extra one of length less than K with one endpoint at x. If possible, this function will return the length of the extra path. Otherwise, the function will return −1.
First, we should call dfs(t) for all children t of x. If any of these return −1, then dfs(x) should also return −1. Otherwise, we have a path of length dfs(t)+1 with one endpoint at x. After doing this, we should pair up as many of the paths whose lengths are in (0,K) as possible. If there is more than one unpaired path remaining after this process, return −1.
Otherwise, return the length of the unpaired path, or zero if none exists. Note that if dfs(x)≠−1, then its return value is equal to the remainder when the number of edges in the subtree corresponding to x is divided by K, which is invariant regardless of how exactly we choose the paths of length K.
For a fixed K, we can check whether it is possible to split the tree into paths of length K in O(N) time, allowing us to solve the problem in O(N⋅(# of divisors of N)). However, several solutions with this complexity ended up receiving TLE on test case 3, where N=83161 and N−1 has 128 divisors. One option is to deal with the star case separately. Another is to write a checker that does not use recursion and is a constant factor faster, demonstrated below.
#include "bits/stdc++.h" using namespace std; void setIO(string s) { ios_base::sync_with_stdio(0); cin.tie(0); freopen((s+".in").c_str(),"r",stdin); freopen((s+".out").c_str(),"w",stdout); } #define f first #define s second const int MOD = 1e9+7; const int MX = 1e5+5; int N,sub[MX]; vector<int> adj[MX], num[MX]; bool bad = 0; void dfs(int a, int b) { sub[a] = 1; for(auto& t: adj[a]) if (t != b) { dfs(t,a); sub[a] += sub[t]; num[a].push_back(sub[t]); } if (sub[a] != N) num[a].push_back(N-sub[a]); } int cur[MX]; // basically unordered map bool ok(int K) { if ((N-1)%K != 0) return 0; for (int i = 0; i < K; ++i) cur[i] = 0; for (int i = 1; i <= N; ++i) { int cnt = 0; for (auto& t: num[i]) { int z = t%K; if (z == 0) continue; if (cur[K-z]) cur[K-z] --, cnt --; else cur[z] ++, cnt ++; } if (cnt) return 0; // paths don't pair up } return 1; } int main() { setIO("deleg"); cin >> N; for (int i = 1; i < N; ++i) { int a,b; cin >> a >> b; adj[a].push_back(b), adj[b].push_back(a); } dfs(1,0); for (int i = 1; i < N; ++i) { if (ok(i)) cout << 1; else cout << 0; } cout << "\n"; }