Submission #956437

#	Time	Username	Problem	Language	Result	Execution time	Memory
956437		caterpillow	Cats or Dogs (JOI18_catdog)	C++17	100 / 100	217 ms	66900 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}); // extra cat case
        dog = clamp({actual.f, actual.s + 1}); // extra dog case
    }
    inline pl cmb(pl a) const {
        if (a.f == a.s) return {a.f + same.f, a.s + same.s};
        else if (a.f > a.s) return {a.s + cat.f, a.s + cat.s};
        else return {a.f + dog.f, a.f + dog.s};
    }
    Node operator+(const Node& b) const {
        return {
            actual + b.actual,
            b.cmb(cat),
            b.cmb(same),
            b.cmb(dog)
        };
    }
    void operator+=(pl b) {
        actual.f += b.f;
        actual.s += b.s;
    }
    void operator-=(pl b) {
        actual.f -= b.f;
        actual.s -= b.s;
    }
};

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];
    }
};

int n;
vt<int> state;

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, {-1, -1});
        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 = -1;
        for (int v : adj[u]) {
            if (v == adj[u][0]) res = dfs_hld(v, rt);
            else dfs_hld(v, v);
        }
        if (res == -1) leaf[u] = u;
        else leaf[u] = res;
        return leaf[u];
    }
    void gen(int r = 0) {
        t = 0;
        par[r] = -1;
        dfs_sz(r);
        dfs_hld(r, r);
        F0R (u, n) {
            if (root[u] == u) agg[u] = {0, 0};
        }
    }
    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 path
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

*/

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