Binary search, greedy, dfs

Let us binary search on number of operations done. 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 or up.

#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