Submission #1057928

#	Time	Username	Problem	Language	Result	Execution time	Memory
1057928		caterpillow	Duathlon (APIO18_duathlon)	C++17	100 / 100	90 ms	48116 KiB

count_triplets

#include <bits/stdc++.h>

using namespace std;

using db = long double;
using ll = long long;
using pl = pair<ll, ll>;
using pi = pair<int, int>;
#define vt vector
#define f first
#define s second
#define pb push_back
#define all(x) x.begin(), x.end() 
#define size(x) ((int) (x).size())
#define FOR(i, a, b) for (int i = (a); i < (b); i++)
#define ROF(i, a, b) for (int i = (b) - 1; i >= (a); i--)
#define F0R(i, b) FOR (i, 0, b)
#define endl '\n'
const ll INF = 1e18;
const int inf = 1e9;

template<typename... Args> // tuples
ostream& operator<<(ostream& os, tuple<Args...> t) { 
    apply([&](Args... args) { string dlm = "{"; ((os << dlm << args, dlm = ", "), ...); }, t);
    return os << "}";
}

template<typename T, typename V> // pairs
ostream& operator<<(ostream& os, pair<T, V> p) { return os << "{" << p.f << ", " << p.s << "}"; }

template<class T, class = decltype(begin(declval<T>()))> // iterables
typename enable_if<!is_same<T, string>::value, ostream&>::type operator<<(ostream& os, const T& v) {
    string dlm = "{";
    for (auto i : v) os << dlm << i, dlm = ", ";
    return os << "}";  
}

template <typename T, typename... V>
void printer(string pfx, const char *names, T&& head, V&& ...tail) {
    int i = 0;
    while (names[i] && names[i] != ',') i++;
    constexpr bool is_str = is_same_v<decay_t<T>, const char*>;
    if (is_str) cerr << " " << head;
    else cerr << pfx, cerr.write(names, i) << " = " << head; 
    if constexpr (sizeof...(tail)) printer(is_str ? "" : ",", names + i + 1, tail...);
    else cerr << endl;
}

#ifdef LOCAL
#define dbg(...) printer(to_string(__LINE__) + ": ", #__VA_ARGS__, __VA_ARGS__)
#else
#define dbg(x...)
#define cerr if (0) std::cerr
#endif

/*

consider when its *not* possible to find a path between three nodes a, b and c
obv, they need to be in the same connected component
clearly, you can go from a to b, and you can also go from b to c
we just need to account for the condition that you cant use the same node twice

consider the block cut tree
a triple is good as long as a, b, and c belong on the same path, in that order

do subtree dp, counting the # of paths with it as their lca
be very careful, since neighbours in the block cut tree share a node which we must not overcount
when counting for some lca, we exclude the shared node with children

cases:
    - we can take a pair from some subtree (excluding shared) and one from another
    - we can take a node each from two subtrees (excluding shared) and one in the root
    - we can take a pair from a subtree (excluding shared) and one from the root (excluding shared)

    additionally, if the root has sufficient nodes
    - we can take two nodes from the root and some node (excluding shared) in a subtree
    - we can take three nodes from the root

we need to track in our dp:
    - how many unique child nodes there are in the subtree (careful overcounting articulation points)
        - this is size of yourself + dp value of children - # of children
    - how many ordered pairs there are in the subtree that lead to the root EXCLUDING the shared node with the parent
        - sum of pairs of children, and (sum of child sizes excluding shared) * (size of root - 1), and (size - 1) * (size - 2)

sum over all roots

*/

struct BCC {
    int n, t; 
    vt<vt<int>> adj;
    vt<int> tin, low, stk;
    vt<bool> is_art;
    vt<vt<int>> comps;

    void init(int _n) {
        n = _n;
        t = 0;
        adj.resize(n);
        tin = low = vt<int>(n);
        is_art.resize(n);
    }

    void ae(int u, int v) {
        adj[u].pb(v);
        adj[v].pb(u);
    }

    void dfs(int u) {
        tin[u] = low[u] = ++t;
        stk.pb(u);
        for (int v : adj[u]) {
            if (tin[v]) low[u] = min(low[u], tin[v]);
            else {
                dfs(v);
                low[u] = min(low[u], low[v]);
                if (low[v] == tin[u]) {
                    is_art[u] = (u != stk[0]) || tin[v] > tin[u] + 1;
                    comps.pb({u});
                    do {
                        comps.back().pb(stk.back());
                        stk.pop_back();
                    } while (comps.back().back() != v);
                }
            }
        }
    }

    void gen() {
        F0R (i, n) if (!tin[i]) dfs(i);
        for (auto c : comps) dbg(c);
    }
};

int n;
BCC bcc;
vt<vt<int>> adj;
vt<int> id;
vt<ll> sz;
vt<vt<int>> rcomp;

vt<ll> sub, pairs;
vt<bool> seen;
ll ans = 0;

void dfs(int u, int p) {
    seen[u] = true;

    for (int v : adj[u]) if (v != p) dfs(v, u);

    // do the first three cases, since they are shared

    ll available = 0; // # of nodes in children that arent part of the root
    for (int v : adj[u]) {
        if (v == p) continue;
        available += sub[v] - 1; // exclude the shared node
    }

    // case 1
    for (int v : adj[u]) {
        if (v == p) continue;
        ans += pairs[v] * (available - (sub[v] - 1)) * 2;        
    }

    // case 2
    for (int v : adj[u]) {
        if (v == p) continue;
        ans += sz[u] * (sub[v] - 1) * (available - (sub[v] - 1));
    }

    // case 3
    for (int v : adj[u]) {
        if (v == p) continue;
        ans += (sz[u] - 1) * pairs[v] * 2;
    }

    // case 4
    for (int v : adj[u]) {
        if (v == p) continue;
        ans += (sz[u] - 1) * (sub[v] - 1) * 2 + (sz[u] - 1) * (sz[u] - 2) * (sub[v] - 1) * 2;
    }

    // case 5
    ans += sz[u] * (sz[u] - 1) * (sz[u] - 2);

    // calc dp values
    sub[u] = available + sz[u];

    pairs[u] = available * (sz[u] - 1) + (sz[u] - 1) * (sz[u] - 2);
    for (int v : adj[u]) {
        if (v == p) continue;
        pairs[u] += pairs[v];
    }
    dbg("\n", rcomp[u]);
    dbg(ans, sub[u], pairs[u], sz[u]);
}

/*

do subtree dp, counting the # of paths with it as their lca
be very careful, since neighbours in the block cut tree share a node which we must not overcount
when counting for some lca, we exclude the shared node with children

cases:
    - we can take a pair from some subtree (excluding shared) and one from another
    - we can take a node each from two subtrees (excluding shared) and one in the root
    - we can take a pair from a subtree (excluding shared) and one from the root (excluding shared)

    additionally, if the root has sufficient nodes
    - we can take two nodes from the root and some node (excluding shared) in a subtree
        - case 1: includes the shared node
        - case 2: exclues shared node
    - we can take three nodes from the root (any)

we need to track in our dp:
    - how many unique child nodes there are in the subtree (careful overcounting articulation points)
        - this is size of yourself + dp value of children - # of children
    - how many ordered pairs there are in the subtree that lead to the root EXCLUDING the shared node with the parent
        - sum of pairs of children, and (sum of child sizes excluding shared) * (size of root - 1), and (size - 1) * (size - 2)

sum over all roots

*/

main() {
    cin.tie(0)->sync_with_stdio(0);
    
    int _n; cin >> _n;
    bcc.init(_n);
    int m; cin >> m;
    F0R (i, m) {
        int u, v; cin >> u >> v; u--, v--;
        bcc.ae(u, v);
    }
    bcc.gen();

    // build block cut tree
    id.resize(_n);
    F0R (i, _n) {
        if (bcc.is_art[i]) {
            id[i] = n++;
            sz.pb(1);
            adj.pb({});
            rcomp.pb({i});
        }
    }
    for (vt<int>& comp : bcc.comps) {
        adj.pb({});
        sz.pb(size(comp));
        for (int u : comp) {
            if (bcc.is_art[u]) {
                adj[id[u]].pb(n);
                adj[n].pb(id[u]);
            }
        }
        rcomp.pb(comp);
        n++;
    }
    dbg(bcc.is_art);

    // start counting
    dbg(n);
    sub = pairs = vt<ll>(n);

    seen.resize(n);
    F0R (i, n) {
        if (!seen[i]) dfs(i, -1);
    }

    cout << ans << endl;
}

Compilation message (stderr)

count_triplets.cpp:225:1: warning: ISO C++ forbids declaration of 'main' with no type [-Wreturn-type]
  225 | main() {
      | ^~~~

• Subtask 1: 5 / 5
• Subtask 2: 11 / 11
• Subtask 3: 8 / 8
• Subtask 4: 10 / 10
• Subtask 5: 13 / 13
• Subtask 6: 15 / 15
• Subtask 7: 20 / 20
• Subtask 8: 8 / 8
• Subtask 9: 10 / 10