#include <bits/stdc++.h>

using namespace std;

void dfs(int u, int par, vector <long long> &a, vector <int> &B, vector <vector <int>> &graph)
{
    bool ch = 0;
    for(auto &it : graph[u])
    {
        if(it == par)
        {
            continue;
        }
        dfs(it, u, a, B, graph);
        if(B[it])
        {
            ch = 1;
        }
    }
    if(B[u])
    {
        ch = 1;
    }
    a[u] = ch;
}

void dfs1(int u, int par, vector <vector <int>> &graph, vector <long long> &a)
{
    long long ans = 1;
    for(auto &it : graph[u])
    {
        if(it == par)
        {
            continue;
        }
        dfs1(it, u, graph, a);
        ans = max(ans, 1 + a[it]);
    }

    a[u] = ans;
}

void dfs2(int u, int par, vector <vector <int>> &graph, vector <int> &x, vector <long long> &a)
{
    int temp;
    bool ch = 1;
    x.push_back(u);
    for(auto &it : graph[u])
    {
        if(it == par)
        {
            continue;
        }
        if(ch && a[u] - 1 == a[it])
        {
            ch = 0;
            temp = it;
        }
        else
        {
            dfs2(it, u, graph, x, a);
            x.push_back(u);
        }
    }
    if(!ch)
    {
        dfs2(temp, u, graph, x, a);
        x.push_back(u);
    }
}

pair<vector<int>, vector<long long>>
find_rebalancing_strategy(int N,
 vector<int> A,
 vector<int> B,
 vector<int> U,
 vector<int> V)
{
    vector <vector <int>> graph(N);

    for(int i = 0 ; i < N - 1 ; ++i)
    {
        graph[U[i]].push_back(V[i]);
        graph[V[i]].push_back(U[i]);
    }
    int n = N;
    vector <int> x;
    vector <long long> y;
    int k;
    for(int i = 0 ; i < n ; ++i)
    {
        if(A[i] > 0)
        {
            k = i;
        }
    }

    vector <long long> a(n, 0);

    dfs(k, -1, a, B, graph);
    a[k] = 1;

    for(auto &it : graph)
    {
        vector <int> temp;
        for(auto &i : it)
        {
            if(a[i])
            {
                temp.push_back(i);
            }
        }
        it = temp;
    }
    for(int i = 0 ; i < n ; ++i)
    {
        a[i] = 0;
    }

    dfs1(k, -1, graph, a);

    dfs2(k, -1, graph, x, a);

    int ma = *max_element(a.begin(), a.end());
    ma --;

    while(ma --)
    {
        x.pop_back();
    }

    for(auto &it : x)
    {
        y.push_back(B[it] - A[it]);
        A[it] = B[it];
    }
    return {x, y};
}


signed main() {
    pair <vector <int>, vector <long long>> ans = find_rebalancing_strategy(5, {0, 11, 0, 0, 0}, {1, 3, 1, 6, 0}, {0, 1, 1, 2}, {1, 2, 3, 4});
    for(auto &it : ans.first)
    {
        cout << it << ' ';
    }
    cout << endl;
    for(auto &it : ans.second)
    {
        cout << it << ' ';
    }
    cout << endl;
}

Have the first person keep eating the pancakes. The second guy will keep eating until he sees an already finished bag of pancakes, in which case he stops. The process of extracting the allergens is trivial. (Why?)

For subtasks 2, 3, and 4 (and I assume for the rest of the problem), we have that the person who has the $$$p$$$-th flavor as their allergy eats the $$$(p + 1 \bmod n)$$$th flavor to convey what index they are. Let $$$v = p + 1 \bmod n$$$.

Eat the $$$v$$$th flavor $$$n-r$$$ times (these are in separate bags). So, if $$$n=5$$$, I am in room 0, and get flavor 4, I eat the 0th flavor 5 times.

Then, we can use the frequency of how many times a flavor was eaten to retrieve the allergens.

Let $$$v = p + 1 \bmod n$$$.

Since every block is a permutation, let's split our info into $$$n$$$ blocks and stagger what each person get. We can create an algorithm that guarantees the following:

• the person in the first room eats one slice
• the person in the second room eats two slices
• the person in the third room eats three slices
• ...

The general algorithm goes as follows:

1. Initialize $$$x = 1$$$. This indicates "the last block of size $$$n$$$" that the person in the room will take from. So, once $$$b \ge xn$$$, we stop eating.
2. If there are 0 slices at some $$$b$$$, update $$$x = \lfloor \frac b n \rfloor+ 1$$$.
3. If $$$f = v$$$, eat it iff $$$b \lt xn$$$, otherwise don't eat it.

Again, let $$$v = p + 1 \bmod n$$$.

The case of $$$n = 2$$$ is dealt with using Subtask 1. Hence, assume $$$n \ge 3$$$.

We reserve the first $$$n+1$$$ bags to denote how many people have passed. With this, it is always possible to take at least one pancake. Therefore, we can correctly extract $$$r$$$. We use that to eat $$$n-r$$$ pancakes of flavor $$$v$$$...

Unless if flavor $$$v$$$ itself appeared $$$n$$$ times in the first $$$n+1$$$ bags, as now we can't report our room! In that case, we realize that the flavor $$$v+1 \bmod n = p+2 \bmod n$$$ can be used! We can do the following if all of the flavor $$$v$$$ pancakes get eaten after the first $$$n+1$$$ rooms:

1. Update $$$v$$$ to $$$v+1 \bmod n$$$.
2. Now, since flavor $$$v$$$ appears at least 2 times, we eat the second appearance of flavor $$$v$$$ to signify our room number.
3. Same deal, eat $$$n-r$$$ pancakes of the 2nd appearance of flavor $$$v$$$.

Implementation for allergen retrieval is a bit tedious, but it is doable.

Each guest in a room looks at the first occurrence of a flavor. If it is eaten, that means this flavor is used by a previous guest, so ignore. If it is not eaten, then use this flavor, i.e. eat this first occurrence. Let this color be $$$q[i]$$$. He can convey his forbidden flavor $$$p[i]$$$ by eating the $$$p[i]$$$-th occurrence of $$$q[i]$$$. The guest outside can recover the array $$$q[i]$$$ just by looking at the sequence of flavors, and recover $$$p[i]$$$ by looking by the eaten bags.

For example, if the sequence of the flavors is $$$2, 1, 3, 3, 1, 2, 1, 0, 2, ....$$$ for $$$n=4$$$, then $$$q=[2,1,3,0]$$$.

The problem arises when the flavor that a guest wants to use is forbidden, i.e. $$$p[i]=q[i]$$$. In this case, take the next unused flavor $$$q[i+1]$$$, and eat the first and the $$$q[i+1]$$$-th bag of flavor $$$q[i+1]$$$. Then when the guest outside recovers $$$q[i]$$$ and $$$p[i]$$$, he finds that $$$q[i+1]=p[i+1]$$$ which should not happen (how can you eat a forbidden flavor?), so he knows that guest $$$i+1$$$ is actually guest $$$i$$$, and guest $$$i$$$ is guest $$$i+1$$$. Therefore, perform the substitution $$$p[i+1]:=p[i],p[i]:=q[i].$$$ If this happens to the last guest, then there is nothing he can eat, but that is fine, because when we get that the last flavor is unused we just set $$$p[n-1]:=q[n-1]$$$.

#include "party.h"
#include <bits/stdc++.h>
using namespace std;
namespace {
    int n, k, p, t, a, y; 
    vector<int> c, q, r, u;
} 
void init(int N, int K, int P, int R) {
	n = N, k = K, p = P, t = 0, a = -1, y = 0;
	c.clear(); q.clear(); r.clear(); u.clear();
	c.resize(n); q.resize(n,-1); r.resize(n,-1); u.resize(n,0);
    return;
}
int strategy(int B, int F, int S) {
	if(!S) u[F] = 1;
	c[F]++;
	if(a != -1 && a != F) return 0;
	if(q[F] == -1) {
		q[F] = t, t++;
		if(!u[F]) {
			if(F == p) y = 1;
			else {a = F; return k;}
		}
	}
	if(y == 0 && F == a && c[F]-1 == p || y == 1 && c[F]-1 == a) return k;
	return 0;
}
vector<int> guess(int N, int K, vector<int> f, vector<int> s) {
	vector<int> C(N,0), P(N,-1), Q(N,-1), R(N,-1), U(N); int T = 0;
	for(int i=0;i<N*N;i++) {
		if(U[f[i]] == 0) P[f[i]] = T, Q[T] = f[i], U[f[i]] = 1, T++;
		if(s[i] == 0) R[P[f[i]]] = C[f[i]];
		C[f[i]]++;
	}
	for(int i=1;i<N;i++) if(R[i] == Q[i]) R[i] = R[i-1], R[i-1] = Q[i-1];
	if(R[N-1] == -1) R[N-1] = Q[N-1];
	return R;
}

Let the graph be a line graph with $$$A = [1,1,0,2,0,0]$$$, $$$B = [0,0,3,0,0,1]$$$, so the nodes' values are $$$[1,1,-3,2,0,-1]$$$.

In your construction, $$$G$$$ is $$$0 \rightarrow 1 \rightarrow 2 \leftarrow 3 \rightarrow 4 \rightarrow 5$$$, so your answer is $$$2 \times 5 - 2 = 8$$$. However, a shorter path is $$$X = [0,1,2,3,2,3,4,5]$$$ and $$$Y = [-1,-1,0,-2,3,0,0,1]$$$.

Root at centroid of the tree (using subtree sizes = number of stations in subtree). This ensures every size <= 3 and you get solution with S = 12.

if size = 3, we can just assume it is a child, this will not change the answer since 6 — 3 = 3. Thus, we can keep a single state, giving S = 10.

Since B is large enough, we can use depth modulo 2 instead of 3. To determine if the edge is parent or child just count the number of edges in both modulo classes. Since count of child edges is always greater, we can identify if it is child edge or parent edge. This gives S = 7.

For nodes of degree $$$2$$$, the sum of labels on the edges is always $$$k$$$, so we need to find a way to differentiate between the two types of nodes (all edges pointing outwards or all edges pointing inwards).

Consider adding $$$k+1$$$ to the labels of edges at depth congruent to $$$0$$$ or $$$1 \pmod 4$$$. Notice that a degree $$$2$$$ node has its adjacent edges pointing towards it if and only if it has an even number of adjacent edges that have label $$$\geq k+1$$$. Now that we can figure out the parity of the depth of a degree $$$2$$$ node based on the labels, we can proceed with the same solution as subtasks 3 and 4.