Binary search, greedy, dfs

Let us binary search on number of operations done. If we design a check function for "Can we do $$$x$$$ moves? It must satisfy a $$$true, ..., true, false, ..., false$$$ scenario. So we can determine the number of moves that each node must move, how to check if it is moving it around possible?

Greedy algorithm used:

Find the maximum number of steps it can move down with, if it can be moved down $$$X$$$ steps, move down. If there is no way to move $$$X$$$ steps down, move up and consider it later.

When would this algorithm stop? It is when some of our nodes are unable to move down and up. This is because when we use a dfs, it ensures that every subtree below it has already satisfied the conditions, we only need to consider current node.

#include<bits/stdc++.h>
using namespace std;

using i64 = long long;

void solve() {
	int N; cin >> N;
	vector<vector<int>> g(N + 1);
	for (int i = 0; i < N - 1; i++) {
		int u, v; cin >> u >> v;
		g[u].push_back(v); g[v].push_back(u);
	}
	int K; cin >> K;
	vector<int> step(N + 1), f(N + 1), col(N + 1), t(K + 1);
	for (int i = 1; i <= K; i++)
		cin >> t[i];

	function<bool(int, int)> dfs = [&](int u, int fa) {
		for (auto& v : g[u]) {
			if (v == fa) continue; //don't go back up first
			if (dfs(v, u) == 0) return 0;//if some node cannot work return false
			if (!col[v]) f[u] = max(f[u], f[v] + 1); //if it is white, f[u]=max(f[u],f[v]+1) because we can use one more operation done
			//just record maximum
		}
		if (step[u] <= f[u]) {//if we can move down, just move down
			step[u] = 0;//remove that node's steps 
			return 1; //return possible
		}
		if (fa == 0 || col[fa]) return 0; //since my code roots at 1, I set its father to be 0
		//if my father is already colored black or it is the root, return false
		step[fa] = step[u] - 1; step[u] = 0; col[fa] = 1;//or else move up, now your steps would - 1
		return 1;
	};

	auto check = [&](int x) {
		for (int i = 0; i <= N; i++) f[i] = step[i] = col[i] = 0; //clear everything
		for (int i = 1; i <= K; i++) {
			int u = t[i];
			step[u] = (x / K);
			if (i <= x % K) step[u]++;//counting step[i] for each a[i]
			col[u] = 1;
		}
		return dfs(1, 0);
	};
	int l = 0, r = N, ans = -1;
	while (l <= r) {
		int mid = (l + r) >> 1;
		if (check(mid)) { 
			ans = mid; 
			l = mid + 1;
		}
		else {
			r = mid - 1;
		}
	}
	cout << ans << '\n';
}

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

	int T; cin >> T;
	while (T--) {
		solve();
	}
}

Interesting problem with a very cool implementation.