Submission #600739

#	Time	Username	Problem	Language	Result	Execution time	Memory
600739		AA_Surely	Construction of Highway (JOI18_construction)	C++14	100 / 100	1237 ms	22148 KiB

construction

#include <bits/stdc++.h>

#define FOR(i, x, n) for(int i = x; i < n; i++)
#define F0R(i, n) FOR(i, 0, n)
#define ROF(i, x, n) for(int i = n - 1; i >= x; i--)
#define R0F(i, n) ROF(i, 0, n)

#define WTF cout << "WTF" << endl

#define IOS ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0)
#define F first
#define S second
#define PB push_back

#define ALL(x) x.begin(), x.end()
#define RALL(x) x.rbegin(), x.rend()

using namespace std;

typedef long long LL;

typedef pair<int, int> PII;
typedef pair<LL, LL> PLL;

typedef vector<int> VI;
typedef vector<LL> VLL;
typedef vector<PII> VPII;
typedef vector<PLL> VPLL;

#define lc now << 1
#define rc now << 1 | 1

const int N = 1e5 + 5;

int n;
int cs[N], sz[N], t[N], cnt;
int head[N], height[N], par[N];
int order[N];
PII eset[N];
VI adj[N], comp;

struct Fenwick {
    int fen[N];
    VI changed;

    void add(int id, int val) {
        for(; id < N; id += id & -id) {
            fen[id] += val;
            changed.PB(id);
        }
        return;
    }

    int get(int id) {
        int ret = 0;
        for(; id; id -= id & -id) ret += fen[id];
        return ret;
    }

    void clear() {
        for(const int &on : changed) fen[on] = 0;
        changed.clear();
    }
} sum;

struct Segment {
    bool tree[N << 2];
    int lz[N << 2], cnt = 0;

    void init() {
        memset(lz, -1, sizeof lz);
        return;
    }

    void shift(int now) {
        if (lz[now] == -1) return;
        
        tree[lc] = lz[lc] = lz[now];
        tree[rc] = lz[rc] = lz[now];

        lz[now] = -1;
        return;
    }

    void change(int lq, int rq, bool val, int now = 1, int ls = 0, int rs = n - 1) {
        if (rq < lq || rq < ls || rs < lq) return;
        if (lq <= ls && rs <= rq) {
            tree[now] = lz[now] = val;
            return;
        }

        int mid = (ls + rs) >> 1;
        shift(now);
        change(lq, rq, val, lc, ls, mid);
        change(lq, rq, val, rc, mid + 1, rs);
        
        tree[now] = tree[lc] | tree[rc];
        return;
    }

    int rightest1(int q, int now = 1, int ls = 0, int rs = n - 1) {
        if (ls >= q || !tree[now]) return -1;
        if (ls == rs) return ls;

        int mid = (ls + rs) >> 1;
        shift(now);
        int case1 = rightest1(q, rc, mid + 1, rs);
        return (case1 == -1 ? rightest1(q, lc, ls, mid) : case1);
    }
} hldq;

struct MaxSeg {
    PII tree[N << 2];
    int cnt = 0;

    void init() {
        fill(tree, tree + (N << 2), PII(-1, -1));
        return;
    }

    void cseg(int lq, int rq, PII val, int now = 1, int ls = 0, int rs = n - 1) {
        if (rq < lq || rq < ls || rs < lq) return;
        if (lq <= ls && rs <= rq) {
            tree[now] = val;
            return;
        }

        int mid = (ls + rs) >> 1;
        cseg(lq, rq, val, lc, ls, mid);
        cseg(lq, rq, val, rc, mid + 1, rs);

        return;
    }

    void change(int lq, int rq, int val) {
        cseg(lq, rq, PII(cnt, val));
        cnt++;
    }
    
    PII gseg(int id, int now = 1, int ls = 0, int rs = n - 1) {
        if (ls == rs) return tree[now];
        int mid = (ls + rs) >> 1;
        if (id <= mid) return max(tree[now], gseg(id, lc, ls, mid));
        return max(tree[now], gseg(id, rc, mid + 1, rs));
    }

    int get(int id) {
        return gseg(id).S;
    }
} color;

void init() {
    cin >> n;
    comp.resize(n);
    F0R(i, n) {
        cin >> cs[i];
        comp[i] = cs[i];
    }

    par[0] = -1;

    sort(ALL(comp)); 
    comp.resize(unique(ALL(comp)) - comp.begin());
    F0R(i, n) cs[i] = lower_bound(ALL(comp), cs[i]) - comp.begin() + 1;

    F0R(i, n - 1) {
        int u, v; cin >> u >> v;
        adj[--u].PB(--v);
        eset[i] = {u, v};

        par[v] = u;
    }

    return;
}

int preD(int now) {
    sz[now] = 1;
    for(int on : adj[now]) {
        height[on] = height[now] + 1;
        sz[now] += preD(on);
    }

    return sz[now];
}

void dfs(int now) {
    order[cnt] = now;
    t[now] = cnt++;

    sort(ALL(adj[now]), [](const int &a, const int &b) {return sz[a] > sz[b];});
    if (!adj[now].empty()) head[ adj[now][0] ] = head[now];

    for(int on : adj[now]) 
        dfs(on);

    return;
}

LL get(int v) {
    LL ret = 0;
    sum.clear();

    while(v != -1) {
        int pl = hldq.rightest1(t[v]);
        if (pl >= t[ head[v] ]) {
            pl = order[pl];

            LL us = height[v] - height[pl];
            int num = color.get(t[v]);

            ret += sum.get(num - 1) * us;
            sum.add(num, us);

            v = pl;
        }

        else {
            LL us = height[v] - height[ head[v] ] + 1;
            int num = color.get(t[v]);
            
            ret += sum.get(num - 1) * us;
            sum.add(num, us);

            v = par[ head[v] ];
        }
    }

    return ret;
}

void change(int v, int col) {
    while(v != -1) {
        hldq.change(t[v], t[v], 1);
        //cout << "changing " << t[head[v]] << ' ' << t[v] << " to " << x
        color.change(t[ head[v] ], t[v], col);

        if (v == head[v]) {
            v = par[v];
            continue;
        }

        v = par[v];
        hldq.change(t[ head[v] ], t[v], 0);
        v = par[ head[v] ];
    }

    return;
}

int main() {
    IOS;

    init();

    /*color.change(0, 0, 1);
    color.change(0, 0, 2);
    cout << color.get(0) << endl;*/

    preD(0);
    iota(head, head + n, 0);
    dfs(0);

    color.init();
    color.change(0, 0, cs[0]);
    hldq.init();

    //F0R(i, n) cout << order[i] + 1 << ' '; cout << endl;

    F0R(i, n - 1) {
        cout << get(eset[i].F) << endl;
        change(eset[i].S, cs[ eset[i].S ]);

        //cout << "after edge " << i + 1 << " colors are ";
        //F0R(i, n) cout << color.get(i) << ' '; cout << endl;
    }

    return 0;
}

Compilation message (stderr)

construction.cpp: In member function 'void Segment::shift(int)':
construction.cpp:78:27: warning: suggest parentheses around assignment used as truth value [-Wparentheses]
   78 |         tree[lc] = lz[lc] = lz[now];
      |                    ~~~~~~~^~~~~~~~~
construction.cpp:79:27: warning: suggest parentheses around assignment used as truth value [-Wparentheses]
   79 |         tree[rc] = lz[rc] = lz[now];
      |                    ~~~~~~~^~~~~~~~~
construction.cpp: In member function 'void Segment::change(int, int, bool, int, int, int)':
construction.cpp:88:33: warning: suggest parentheses around assignment used as truth value [-Wparentheses]
   88 |             tree[now] = lz[now] = val;
      |                         ~~~~~~~~^~~~~

• Subtask 1: 7 / 7
• Subtask 2: 9 / 9
• Subtask 3: 84 / 84