Submission #956440

#	Time	Username	Problem	Language	Result	Execution time	Memory
956440		caterpillow	Cats or Dogs (JOI18_catdog)	C++17	100 / 100	244 ms	68472 KiB

catdog

#include <bits/stdc++.h>
#include "catdog.h"

using namespace std;

using ll = long long;
using pl = pair<ll, ll>;
#define vt vector
#define f first
#define s second
#define all(x) x.begin(), x.end() 
#define pb push_back
#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'
#define debug(x) do{auto _x = x; cerr << #x << " = " << _x << endl;} while(0)
const ll INF = 1e9;

pl clamp(pl a) {
    return {min(a.f, a.s + 1), min(a.f + 1, a.s)};
}

pl operator+(const pl& a, const pl& b) {
    return {a.f + b.f, a.s + b.s};
}

pl operator-(const pl& a, const pl& b) {
    return {a.f - b.f, a.s - b.s};
}

struct Node {
    pl actual, cat, same, dog; // actual sum, extra cat, same, extra dog
    void build() {
        same = clamp(actual); // same case
        cat = clamp({actual.f + 1, actual.s}); // assume an extra cat
        dog = clamp({actual.f, actual.s + 1}); // assume an extra dog 
    }
    inline pl cmb(pl a) const {
        if (a.f == a.s) return a + same; // same case
        else if (a.f > a.s) return pl {a.s, a.s} + cat; // extra cat case
        else return pl {a.f, a.f} + dog; // extra dog case
    }
    Node operator+(const Node& b) const {
        return {
            actual + b.actual,
            b.cmb(cat),
            b.cmb(same),
            b.cmb(dog)
        };
    }
    void operator+=(pl b) {
        actual = actual + b;
    }
    void operator-=(pl b) {
        actual = actual - b;
    }
};

const Node NID = {
    {0, 0},
    {1, 0},
    {0, 0},
    {0, 1}
};

// cursed reverse segtree
struct SegTree {
    int n;
    vt<Node> seg; 
    void init(int _n) {
        for (n = 1; n < _n; n *= 2);
        seg.resize(2 * n, NID);
    }
    Node query(int l, int r) {
        Node lhs = NID, rhs = NID;
        for (l += n, r += n + 1; l < r; l /= 2, r /= 2) {
            if (l & 1) lhs = seg[l++] + lhs;
            if (r & 1) rhs = rhs + seg[--r];
        }
        return rhs + lhs;
    }
    void pull(int i) {
        i += n;
        while (i > 1) i /= 2, seg[i] = seg[2 * i + 1] + seg[2 * i];
    }
    void upd(int i, pl val) {
        seg[i += n] += val;
        seg[i].build();
        pull(i - n);
    }
    Node& operator[](int i) {
        return seg[i + n];
    }
};

struct HLD {
    int n, t;
    vt<vt<int>> adj;
    vt<int> pos, root, sz, par, leaf;
    vt<pl> agg;
    SegTree seg;
    void init(int _n) {
        n = _n;
        adj.resize(n);
        agg.resize(n);
        pos = root = sz = par = leaf = vt<int>(n);
        seg.init(n);
    }
    void ae(int u, int v) {
        adj[u].pb(v);
        adj[v].pb(u);
    }
    int dfs_sz(int u) {
        sz[u] = 1;
        for (int& v : adj[u]) {
            par[v] = u;
            adj[v].erase(find(all(adj[v]), u));
            sz[u] += dfs_sz(v);
            if (sz[v] > sz[adj[u][0]]) swap(v, adj[u][0]);
        }
        return sz[u];
    }
    int dfs_hld(int u, int rt) {
        root[u] = rt;
        pos[u] = t++;
        int res = u;
        for (int v : adj[u]) {
            if (v == adj[u][0]) res = dfs_hld(v, rt);
            else dfs_hld(v, v);
        }
        return leaf[u] = res;
    }
    void gen(int r = 0) {
        t = 0;
        par[r] = -1;
        dfs_sz(r);
        dfs_hld(r, r);
    }
    void prop(int u) {
        while (true) {
            u = leaf[u]; // go to bottom of heavy path

            // calc updates subtree aggregate
            int rt = root[u];
            Node res = seg.query(pos[rt], pos[u]);

            // update aggregate and its parent
            int p = par[rt];

            if (p >= 0) seg[pos[p]] -= agg[rt];
            agg[rt] = res.same;
            if (p == -1) return;
            
            seg[pos[p]] += res.same;
            seg[pos[p]].build();
            seg.pull(pos[p]);

            u = p; // go to parent
        }
    }  
    ll update(int u, pl v) { // do stuff to update u
        seg.upd(pos[u], v);
        prop(u);
        return min(agg[0].f, agg[0].s);
    }
} hld;

/*

let dp[u] = {min cost to make u's subtree cat safe, min cost to make u's subtree dog safe}
the difference between them is at most 1

keep track of all the cases (extra cat, same, extra dog)
each node stores the sum of everything EXCEPT its heavy child
heavy path roots store its subtree aggregate

when performing an update, modify the current node
walk up from the BASE of each heavy path up to the root, updating subtree aggregates and their parents

*/

int n;
vt<int> state;

void initialize(int N, vt<int> A, vt<int> B) {
    n = N;
    state.resize(n);
    hld.init(n);
    F0R (i, n - 1) {
        int u = A[i] - 1, v = B[i] - 1;
        hld.ae(u, v);
    }
    hld.gen();
}

int cat(int v) {
    v--;
    state[v] = 1;
    return hld.update(v, {0, INF});
}

int dog(int v) {
    v--;
    state[v] = 2;
    return hld.update(v, {INF, 0});
}

int neighbor(int v) {
    v--;
    pl sus;
    if (state[v] == 1) sus = {0, -INF};
    else sus = {-INF, 0};
    state[v] = 0;
    return hld.update(v, sus);
}

• Subtask 1: 8 / 8
• Subtask 2: 30 / 30
• Subtask 3: 62 / 62