Code : Count Paths

Concise Quadratic DP ↕

Submission

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

class Tree {
  public:
    int n;
    vector<vector<int>> adj;
    vector<long long> cnt, dp;
    vector<int> a;

    // cnt[u] is the number of downward paths starting at u when only the
    // last vertex in the path has black color.

    // dp[u] is the number of paths where the terminal vertices are black,
    // they have LCA u and all internal vertices are white.
    Tree(int n) {
        this->n = n;
        adj.resize(n);
        cnt.resize(n);
        a.resize(n);
        dp.resize(n);
    }

    void dfs(int u, int par, int color) {
        cnt[u] = (a[u] == color);

        long long pref_sum = 0;
        for (auto child : adj[u]) {
            if (child == par) {
                continue;
            }
            dfs(child, u, color);

            // If it is a white vertex, LCA is an internal node.
            // To create the subtree, you can pick one element in the incoming
            // child and the other element in any of the subtree seen so far.
            if (a[u] != color) {
                cnt[u] += cnt[child];
                dp[u] += pref_sum * cnt[child];
            } else {
                dp[u] += cnt[child];
            }
            pref_sum += cnt[child];
        }
    }

    void clear() {
        dp.assign(n, 0);
        cnt.assign(n, 0);
    }
};

void solve() {
    int n;
    cin >> n;
    Tree t(n);
    for (int i = 0; i < n; i++) {
        cin >> t.a[i];
        t.a[i]--;
    }

    for (int i = 0; i < n - 1; i++) {
        int x, y;
        cin >> x >> y;
        x--;
        y--;
        t.adj[x].push_back(y);
        t.adj[y].push_back(x);
    }

    long long res = 0;
    for (int color = 0; color < n; color++) {
        // Calculate answer when terminal vertices are of color 'color'.
        t.dfs(0, -1, color);
        long long cur = 0;
        for (int i = 0; i < n; i++) {
            cur += t.dp[i];
        }
        res += cur;
        t.clear();
    }

    cout << res << endl;
}

int main() {
    int t;
    cin >> t;
    for (int i = 0; i < t; i++) {
        solve();
    }
    return 0;
}

Quadratic Choice DP ↕

Submission

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

class Tree {
  public:
    int n;
    vector<vector<int>> adj;
    vector<long long> cnt, dp;
    vector<int> a;

    // cnt[u] is the number of downward paths starting at u when only the
    // last vertex in the path has black color.

    // dp[u] is the number of paths where the terminal vertices are black,
    // they have LCA u and all internal vertices are white.
    Tree(int n) {
        this->n = n;
        adj.resize(n);
        cnt.resize(n);
        a.resize(n);
        dp.resize(n);
    }

    void dfs(int u, int par, int color) {
        cnt[u] = (a[u] == color);

        vector<long long> choices;
        for (auto child : adj[u]) {
            if (child == par) {
                continue;
            }
            dfs(child, u, color);

            if (a[u] != color) {
                cnt[u] += cnt[child];
            }
            choices.push_back(cnt[child]);
        }

        // If the current vertex is white, it can be LCA only if exactly 2
        // nodes from different subtrees are chosen.
        if (a[u] != color) {
            long long pref_sum = 0;
            for (int i = 0; i < choices.size(); i++) {
                dp[u] += pref_sum * choices[i];
                pref_sum += choices[i];
            }
        }

        // If the current node is black, it means that this is a terminal
        // element. Hence, only one subtree is needed.
        if (a[u] == color) {
            for (int i = 0; i < choices.size(); i++) {
                dp[u] += choices[i];
            }
        }
    }

    void clear() {
        dp.assign(n, 0);
        cnt.assign(n, 0);
    }
};

void solve() {
    int n;
    cin >> n;
    Tree t(n);
    for (int i = 0; i < n; i++) {
        cin >> t.a[i];
        t.a[i]--;
    }

    for (int i = 0; i < n - 1; i++) {
        int x, y;
        cin >> x >> y;
        x--;
        y--;
        t.adj[x].push_back(y);
        t.adj[y].push_back(x);
    }

    long long res = 0;
    for (int color = 0; color < n; color++) {
        // Calculate answer when terminal vertices are of color 'color'.
        t.dfs(0, -1, color);
        long long cur = 0;
        for (int i = 0; i < n; i++) {
            cur += t.dp[i];
        }
        res += cur;
        t.clear();
    }

    cout << res << endl;
}

int main() {
    int t;
    cin >> t;
    for (int i = 0; i < t; i++) {
        solve();
    }
    return 0;
}

Cubic Choice DP ↕

Submission

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

class Tree {
  public:
    int n;
    vector<vector<int>> adj;
    vector<long long> cnt, dp;
    vector<int> a;

    // cnt[u] is the number of downward paths starting at u when only the
    // last vertex in the path has black color.

    // dp[u] is the number of paths where the terminal vertices are black,
    // they have LCA u and all internal vertices are white.
    Tree(int n) {
        this->n = n;
        adj.resize(n);
        cnt.resize(n);
        a.resize(n);
        dp.resize(n);
    }

    void dfs(int u, int par, int color) {
        cnt[u] = (a[u] == color);

        vector<long long> choices;
        for (auto child : adj[u]) {
            if (child == par) {
                continue;
            }
            dfs(child, u, color);

            if (a[u] != color) {
                cnt[u] += cnt[child];
            }
            choices.push_back(cnt[child]);
        }

        // If the current vertex is white, it can be LCA only if exactly 2
        // nodes from different subtrees are chosen.
        if (a[u] != color) {
            for (int i = 0; i < choices.size(); i++) {
                for (int j = i + 1; j < choices.size(); j++) {
                    dp[u] += choices[i] * choices[j];
                }
            }
        }

        // If the current node is black, it means that this is a terminal
        // element. Hence, only one subtree is needed.
        if (a[u] == color) {
            for (int i = 0; i < choices.size(); i++) {
                dp[u] += choices[i];
            }
        }
    }

    void clear() {
        dp.assign(n, 0);
        cnt.assign(n, 0);
    }
};

void solve() {
    int n;
    cin >> n;
    Tree t(n);
    for (int i = 0; i < n; i++) {
        cin >> t.a[i];
        t.a[i]--;
    }

    for (int i = 0; i < n - 1; i++) {
        int x, y;
        cin >> x >> y;
        x--;
        y--;
        t.adj[x].push_back(y);
        t.adj[y].push_back(x);
    }

    long long res = 0;
    for (int color = 0; color < n; color++) {
        // Calculate answer when terminal vertices are of color 'color'.
        t.dfs(0, -1, color);
        long long cur = 0;
        for (int i = 0; i < n; i++) {
            cur += t.dp[i];
        }
        res += cur;
        t.clear();
    }

    cout << res << endl;
}

int main() {
    int t;
    cin >> t;
    for (int i = 0; i < t; i++) {
        solve();
    }
    return 0;
}