1843DUnreviewed1200GPT-5.5 (high)

Apple Tree

Progressive hints first, then the full explanation and implementation when you're ready to cash out.

combinatorics dfs and similar dp math trees

Original problem

Review status

AI-generated and still unreviewed. Double-check the details before internalizing them.

Hints

Progressive nudges

Open only as much as you need to keep the solve alive.

Forget the shaking animation for a second. From a vertex $u$ , where can an apple eventually fall? It can only keep moving downward, never upward.

The final falling vertex must be a leaf in the rooted tree. So for each vertex $u$ , you only care about how many leaves exist inside the subtree of $u$ .

Define $dp[u]$ as the number of possible falling vertices for an apple starting at $u$ . If $u$ has no children, then $dp[u] = 1$ .

For an internal vertex, the apple chooses one child and then behaves independently inside that child's subtree, so:

dp[u] = \sum_{v \text{ child of } u} dp[v]

After one DFS from root $1$ computes all $dp[u]$ , every query $(x,y)$ is just:

\boxed{dp[x] \cdot dp[y]}

Use long long, because this can be around $n^2$ . Yes, the whole problem is basically “count leaves in subtrees.” Sneaky little tree gremlin.

Key observation

An apple starting at vertex $u$ always moves from a vertex to one of its children. So it can never leave the subtree of $u$ , and it must eventually stop at some vertex with no children — a leaf of the rooted tree.

Therefore, the set of possible vertices from which the apple falls is exactly:

\{\text{leaves inside the subtree of } u\}

So the only useful value for each vertex is:

dp[u] = \text{number of leaves in the subtree of } u

DP on the rooted tree

Root the tree at vertex $1$ . For every vertex $u$ :

If $u$ has no children, then it is a leaf, so:

dp[u] = 1

Otherwise, every possible final leaf lies inside exactly one child subtree, so:

dp[u] = \sum_{v \text{ child of } u} dp[v]

This is a standard bottom-up tree DP. One DFS is enough.

Because $n$ can be $2 \cdot 10^5$ , a recursive DFS can be a stack-overflow casino. We avoid that nonsense by building a DFS order iteratively, then processing it in reverse.

Answering queries

For a query $(x, y)$ :

The first apple can fall from any of $dp[x]$ leaves.
The second apple can fall from any of $dp[y]$ leaves.
The pair is ordered: $(a,b)$ .

Thus the answer is simply:

ans = dp[x] \cdot dp[y]

Even if $x = y$ , the apples are still considered separately, so the multiplication is still correct.

Complexity

For each test case:

Building/rooting the tree and computing DP: $O(n)$
Answering all queries: $O(q)$

Total complexity over all test cases:

O\left(\sum n + \sum q\right)

Memory usage:

O(n)

Use long long for answers, since $dp[x] \cdot dp[y]$ can be as large as $n^2$ .

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

using ll = long long;

void setIO() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
}

int main() {
    setIO();

    int T;
    cin >> T;

    while (T--) {
        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);
        }

        vector<int> parent(n + 1, 0), order;
        order.reserve(n);

        vector<int> st;
        st.push_back(1);
        parent[1] = -1;

        while (!st.empty()) {
            int u = st.back();
            st.pop_back();
            order.push_back(u);

            for (int v : g[u]) {
                if (v == parent[u]) continue;
                parent[v] = u;
                st.push_back(v);
            }
        }

        vector<ll> dp(n + 1, 0);

        for (int i = n - 1; i >= 0; i--) {
            int u = order[i];
            bool leaf = true;

            for (int v : g[u]) {
                if (v == parent[u]) continue;
                leaf = false;
                dp[u] += dp[v];
            }

            if (leaf) dp[u] = 1;
        }

        int q;
        cin >> q;

        while (q--) {
            int x, y;
            cin >> x >> y;
            cout << dp[x] * dp[y] << '\n';
        }
    }

    return 0;
}

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

using ll = long long;

void setIO() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
}

int main() {
    setIO();

    int T;
    cin >> T;

    while (T--) {
        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);
        }

        vector<int> parent(n + 1, 0), order;
        order.reserve(n);

        vector<int> st;
        st.push_back(1);
        parent[1] = -1;

        while (!st.empty()) {
            int u = st.back();
            st.pop_back();
            order.push_back(u);

            for (int v : g[u]) {
                if (v == parent[u]) continue;
                parent[v] = u;
                st.push_back(v);
            }
        }

        vector<ll> dp(n + 1, 0);

        for (int i = n - 1; i >= 0; i--) {
            int u = order[i];
            bool leaf = true;

            for (int v : g[u]) {
                if (v == parent[u]) continue;
                leaf = false;
                dp[u] += dp[v];
            }

            if (leaf) dp[u] = 1;
        }

        int q;
        cin >> q;

        while (q--) {
            int x, y;
            cin >> x >> y;
            cout << dp[x] * dp[y] << '\n';
        }
    }

    return 0;
}

View on Codeforces Back to the list