Official

G - Delivery on Tree Editorial by en_translator


First of all, the path of the basket is limited to one that passes through every edge exactly once and eventually returns to vertex \(1\). So, determining a path is equivalent to determining for each vertex with two children their order to visit.

We consider conditions imposed to an operation sequence that delivers the balls correctly. First, we realize the following:

  • (Constraint A) for each vertex \(i\) with two children, exactly one of the following is imposed: “the left child must be visited first,” “the right child must be visited first,” or “the order does not matter.”

Specifically, if the lowest common ancestor \(\text{lca}(S_j,T_j)\) of \(S_j\) and \(T_j\) coincide neither with \(S_j\) nor \(T_j\), such a constraint is imposed on vertex \(\text{lca}(S_j,T_j)\). If the constraints conflict, then the answer is obviously \(0\).

Next, when the order to put the balls into our out of the basket is optimal, the following values are determined:

  • (Constraint B) for each vertex \(i\) and for each \(c\in \{\)its left child\(,\)its right child\(,\)its parent\(\}\), “the number of balls to put into the basket right before moving from \(i\) to \(c\)” and “the number of balls to put into the basket right after moving from \(c\) to \(i\).”

For example, if \(S_j\) is an ancestor of \(T_j\), and \(S_j\)’s right child is on the \(S_i-T_j\) path, then it is optimal to put ball \(j\) into the basket right before moving from \(S_j\) to its right child, and out of the basket right after moving from \(T_j\)’s parent to \(T_j\).

All that left is to count the number of paths among those under constraint A such that operations complying constraint B never lets the basket have more than \(K\) balls. This can be found by the following DP (Dynamic Programming):

  • \(dp[i][j]\): when visiting vertex \(i\) for the first time with \(j\) balls in the basket,
    • how many paths within the subtree \(i\) satisfy the constraints?
    • after visiting all the vertices in subtree \(i\) and returning to \(i\)’s parent, how many balls are there in the basket? (Note that, if we let \(j'\) be this value, \(j'-j\) is a constant independent of \(j\).)

The transition itself is natural and there is nothing special to mention, but in fact the main part of this problem is implementation. The following sample code adopts memorized recursion to avoid complexities occurring when trying to evaluate the DP value for an impossible pair \((i,j)\). Also, to avoid complex conditional branches depending on the number of children, it exploits exhaustive search using permutations.

Sample code (C++):

#include<bits/stdc++.h>
#include<atcoder/modint>

using namespace std;
using namespace atcoder;

using mint = modint998244353;

vector<vector<int>> par;
vector<int> dep;

void lca_init() {
    int n = par[0].size();
    for (int k = 0; k < 19; k++) {
        for (int i = 0; i < n; i++) {
            if (par[k][i] == -1) continue;
            par[k + 1][i] = par[k][par[k][i]];
        }
    }
    dep.resize(n);
    for (int i = 1; i < n; i++) {
        dep[i] = dep[par[0][i]] + 1;
    }
}

int la(int u, int d) {
    for (int k = 0; k < 20; k++) {
        if (d >> k & 1) u = par[k][u];
    }
    return u;
}

int lca(int u, int v) {
    if (dep[u] > dep[v]) swap(u, v);
    v = la(v, dep[v] - dep[u]);
    if (u == v) return u;
    for (int k = 19; k >= 0; k--) {
        if (par[k][u] != par[k][v]) {
            u = par[k][u];
            v = par[k][v];
        }
    }
    return par[0][u];
}

int main() {
    int n, m, k;
    cin >> n >> m >> k;
    vector<vector<int>> ch(n);
    par = vector(20, vector<int>(n, -1));
    for (int i = 1; i < n; i++) {
        int p;
        cin >> p;
        ch[--p].push_back(i);
        par[0][i] = p;
    }
    lca_init();
    vector in(n, vector<int>(3)); // 0 : left child  1 : right child  2 : parent
    vector out(n, vector<int>(3));
    vector<int> first(n, -1);
    for (int i = 0; i < m; i++) {
        int s, t;
        cin >> s >> t;
        --s, --t;
        int l = lca(s, t);
        int cs = -1, ct = -1;
        if (l != s) {
            int now = la(s, dep[s] - dep[l] - 1);
            for (int j = 0; j < (int) ch[l].size(); j++) {
                if (ch[l][j] == now) cs = j;
            }
        }
        if (l != t) {
            int now = la(t, dep[t] - dep[l] - 1);
            for (int j = 0; j < (int) ch[l].size(); j++) {
                if (ch[l][j] == now) ct = j;
            }
        }
        if (l == s) {
            ++in[s][ct];
            ++out[t][2];
        } else if (l == t) {
            ++in[s][2];
            ++out[t][cs];
        } else {
            ++in[s][2];
            ++out[t][2];
            if (first[l] == ct) {
                cout << 0 << endl;
                return 0;
            }
            first[l] = cs;
        }
    }
    vector dp(n, vector<pair<int, mint>>(k + 1));
    vector seen(n, vector<bool>(k + 1));
    auto f = [&](auto &f, int i, int j) -> pair<int, mint> {
        if (seen[i][j]) return dp[i][j];
        seen[i][j] = true;
        dp[i][j] = {-1, 0};
        vector<int> ord(ch[i].size());
        iota(ord.begin(), ord.end(), 0);
        do {
            if (!ord.empty() and first[i] == 1 - ord[0]) continue;
            int nj = j;
            mint now = 1;
            nj -= out[i][2];
            for (int c: ord) {
                nj += in[i][c];
                if (nj > k) {
                    nj = -1;
                    break;
                }
                mint m;
                tie(nj, m) = f(f, ch[i][c], nj);
                if (nj == -1) break;
                now *= m;
                nj -= out[i][c];
            }
            if (nj == -1) continue;
            nj += in[i][2];
            if (nj > k) continue;
            dp[i][j].first = nj;
            dp[i][j].second += now;
        } while (next_permutation(ord.begin(), ord.end()));
        return dp[i][j];
    };
    cout << f(f, 0, 0).second.val() << endl;
}

posted:
last update: