BJOI2017 树的难题

#include <cstdio>
#include <cstring>

#include <algorithm>

#define LOG(FMT...) // fprintf(stderr, "[LOG]"FMT)

using namespace std;

typedef long long ll;

struct node {
    int l, r;
    ll d;
    bool flag;
    node *ls, *rs;

    void push_down();
    ll get_d() const;
    ll query(int l, int r);
};

struct edge {
    int v, w;
    edge* next;
};

const int N = 200010;
const ll INF = 1LL << 50;

int n, m, bl, br, ql, qr, ocnt;
ll ans = -INF;
node *seg1, *seg2;
bool vis[N], cut[N];
int c[N], q[N], sum[N], mex[N], ord[N], dep[N], w[N];
ll sc[N];
pair<int, int> son[N];
edge* g[N];

void add_edge(int u, int v, int w);
void dfs_order(int u);
void change(node* p, int x, ll d);
void bfs(int u);
node* build(int l, int r);

int main() {
    scanf("%d%d%d%d", &n, &m, &bl, &br);
    for (int i = 1; i <= m; ++i)
        scanf("%d", &c[i]);
    for (int rep = 1; rep < n; ++rep) {
        int u, v, w;
        scanf("%d%d%d", &u, &v, &w);
        add_edge(u, v, w);
        add_edge(v, u, w);
    }
    dfs_order(1);
    seg1 = build(1, n);
    seg2 = build(1, n);
    memset(cut, 0, sizeof(cut));
    for (int i = 1; i <= n; ++i) {
        int u = ord[i], scnt = 0;
        LOG("%d\n", u);
        cut[u] = true;
        seg1->flag = true;
        seg2->flag = true;
        for (edge* p = g[u]; p; p = p->next)
            if (!cut[p->v]) {
                w[p->v] = p->w;
                sc[p->v] = c[p->w];
                LOG("->%d %d %d\n", p->v, p->w, c[p->w]);
                son[++scnt] = make_pair(p->w, p->v);
            }
        sort(son + 1, son + scnt + 1);
        for (int i = 1; i <= scnt; ++i) {
            int u = son[i].second;
            dep[u] = 1;
            bfs(u);
            for (int j = 0; j < qr; ++j) {
                int v = q[j];
                LOG("SCORE %d: %lld\n", v, sc[v]);
                if (dep[v] > br)
                    break;
                if (dep[v] >= bl)
                    ans = max(ans, sc[v]);
                if (dep[v] < br) {
                    int l = max(bl - dep[v], 1), r = br - dep[v];
                    ans = max(ans, sc[v] + max(seg1->query(l, r), seg2->query(l, r)));
                }
                for (edge* p = g[v]; p; p = p->next)
                    if (!cut[p->v] && dep[v] < dep[p->v]) {
                        w[p->v] = p->w;
                        sc[p->v] = sc[v] + c[p->w];
                        if (w[v] == p->w)
                            sc[p->v] -= c[p->w];
                    }
            }
            for (int j = 0; j < qr; ++j)
                change(seg2, dep[q[j]], sc[q[j]] - c[w[u]]);
            if (son[i].first != son[i + 1].first) {
                for (int j = i; son[i].first == son[j].first; --j) {
                    bfs(son[j].second);
                    for (int k = 0; k < qr; ++k)
                        change(seg1, dep[q[k]], sc[q[k]]);
                }
                seg2->flag = true;
            }
        }
    }
    printf("%lld\n", ans);
    return 0;
}

void bfs(int u) {
    ql = 0;
    qr = 1;
    q[0] = u;
    while (ql < qr) {
        u = q[ql++];
        vis[u] = true;
        for (edge* p = g[u]; p; p = p->next)
            if (!vis[p->v] && !cut[p->v]) {
                dep[p->v] = dep[u] + 1;
                q[qr++] = p->v;
            }
    }
    for (int i = 0; i < qr; ++i)
        vis[q[i]] = false;
}

node* build(int l, int r) {
    static node pool[N * 4];
    static node* top = pool;
    node* p = top;
    ++top;
    p->l = l;
    p->r = r;
    if (l == r)
        return p;
    int mid = (l + r) / 2;
    p->ls = build(l, mid);
    p->rs = build(mid + 1, r);
    return p;
}

void change(node* p, int x, ll d) {
    LOG("CHANGEING %lld\n", d);
    while (p->l != p->r) {
        p->push_down();
        p->d = max(p->d, d);
        if (x <= p->ls->r)
            p = p->ls;
        else
            p = p->rs;
    }
    p->push_down();
    p->d = max(p->d, d);
}

void node::push_down() {
    if (flag) {
        flag = false;
        d = -INF;
        if (ls) {
            ls->flag = true;
            rs->flag = true;
        }
    }
}

ll node::get_d() const {
    if (flag)
        return -INF;
    return d;
}

ll node::query(int l, int r) {
    if (this->l == l && this->r == r)
        return get_d();
    push_down();
    if (r <= ls->r)
        return ls->query(l, r);
    if (l >= rs->l)
        return rs->query(l, r);
    return max(ls->query(l, ls->r), rs->query(rs->l, r));
}

void add_edge(int u, int v, int w) {
    static edge pool[N * 2];
    static edge* p = pool;
    p->v = v;
    p->w = w;
    p->next = g[u];
    g[u] = p;
    ++p;
}

void dfs_order(int u) {
    bfs(u);
    int myn = 1 << 20;
    for (int i = qr - 1; i >= 0; --i) {
        int v = q[i];
        sum[v] = 1;
        mex[v] = 0;
        for (edge* p = g[v]; p; p = p->next)
            if (dep[v] < dep[p->v] && !cut[p->v]) {
                sum[v] += sum[p->v];
                mex[v] = max(mex[v], sum[p->v]);
            }
        mex[v] = max(mex[v], qr - sum[v]);
        if (mex[v] < myn) {
            myn = mex[v];
            u = v;
        }
    }
    cut[u] = true;
    ord[++ocnt] = u;
    for (edge* p = g[u]; p; p = p->next)
        if (!cut[p->v])
            dfs_order(p->v);
}