Submission #947947

#	Time	Username	Problem	Language	Result	Execution time	Memory
947947		Nhoksocqt1	Closing Time (IOI23_closing)	C++17	100 / 100	375 ms	64256 KiB

closing

#ifndef Nhoksocqt1
    #include "closing.h"
#endif // Nhoksocqt1
#include<bits/stdc++.h>
using namespace std;

#define inf 0x3f3f3f3f
#define sz(x) int((x).size())
#define fi first
#define se second
typedef long long ll;
typedef pair<int, int> ii;

template<class X, class Y>
	inline bool maximize(X &x, const Y &y) {return (x < y ? x = y, 1 : 0);}
template<class X, class Y>
	inline bool minimize(X &x, const Y &y) {return (x > y ? x = y, 1 : 0);}

mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());
int Random(int l, int r) {
    return uniform_int_distribution<int>(l, r)(rng);
}

const bool isMultiTest = 0;
const int MAXN = 200005;

vector<ii> adj[MAXN];
ll fen_dist[3 * MAXN], fen_l[3 * MAXN], fen_h[3 * MAXN];
ll dist_l[MAXN], dist_h[MAXN], maxSum;
int pa[MAXN], numNode, x_node, y_node;
bool inPath[MAXN];

void dfs_init(int u, int p, ll dist_from_root, ll dist[], int par[]) {
    par[u] = p;
    dist[u] = dist_from_root;
    for (int it = 0; it < sz(adj[u]); ++it) {
        int v(adj[u][it].fi), w(adj[u][it].se);
        if(v != par[u])
            dfs_init(v, u, dist_from_root + w, dist, par);
    }
}

typedef pair<ll, int> pli;
struct Data {
    ll dif, val_l;
    int node;
};

vector<Data> sorted_diff;
vector<pli> sorted_val;

void modify(ll fen[], int i, ll v) {
    for (; i <= 3 * numNode; i += i & -i)
        fen[i] += v;
}

ll get(ll fen[], int i) {
    ll res(0);
    for (; i > 0; i -= i & -i)
        res += fen[i];

    return res;
}

int walkOnFen(ll fen[], ll maxSum) {
    int i(0);
    for (int j = 31 - __builtin_clz(3 * numNode); j >= 0; --j) {
        if(i + (1 << j) <= 3 * numNode && maxSum >= fen[i + (1 << j)]) {
            i += (1 << j);
            maxSum -= fen[i];
        }
    }

    return i;
}

int getSumOptLeft(ll maxSum) {
    int idx = walkOnFen(fen_dist, maxSum);

    ll max_sum = get(fen_dist, idx);
    int get_l = get(fen_l, idx), get_h = get(fen_h, idx);
    int max_cnt = get_l + get_h;

    int id_h_now = walkOnFen(fen_h, get_h - 1);
    ll dist_h_now = sorted_val[id_h_now].fi;

    if(get_h % 2 == 0)
        return max_cnt;

    // replace highest q with next lowest p
    if(get_l < get(fen_l, 3 * numNode)) {
        int id_l_now = walkOnFen(fen_l, get_l);
        ll dist_l_now = sorted_val[id_l_now].fi;
        if(max_sum - dist_h_now + dist_l_now <= maxSum)
            return max_cnt;
    }

    // replace highest p with highest q
    if(get_l > 0) {
        int id_l_now = walkOnFen(fen_l, get_l - 1);
        ll dist_l_now = sorted_val[id_l_now].fi;
        if(max_sum - dist_l_now + dist_h_now <= maxSum)
            return max_cnt;
    }

    return max_cnt - 1;
}

int calcWithCommon(void) {
    sorted_diff.clear();
    for (int i = 0; i < numNode; ++i)
        sorted_diff.push_back({dist_h[i] - dist_l[i], dist_l[i], i});

    sort(sorted_diff.begin(), sorted_diff.end(), [](const Data &a, const Data &b) {
        return (a.dif != b.dif) ? a.dif < b.dif : a.val_l < b.val_l;
    });

    sorted_val.clear();
    for (int i = 0; i < numNode; ++i) {
        sorted_val.push_back(pli(2 * dist_l[i], i));
        sorted_val.push_back(pli(dist_h[i], -2 * i - 2));
        sorted_val.push_back(pli(dist_h[i], -2 * i - 1));
    }

    sort(sorted_val.begin(), sorted_val.end());
    for (int i = 1; i <= 3 * numNode; ++i)
        fen_dist[i] = fen_l[i] = fen_h[i] = 0;

    ll sum_in_path(0);
    int cnt_in_path(0), res(0);
    for (int u = 0; u < numNode; ++u) {
        if(inPath[u]) {
            sum_in_path += dist_l[u];
            ++cnt_in_path;
            continue;
        }

        int pos = upper_bound(sorted_val.begin(), sorted_val.end(), pli(2 * dist_l[u], u)) - sorted_val.begin();
        modify(fen_dist, pos, 2 * dist_l[u]);
        modify(fen_l, pos, +1);
    }

    if(sum_in_path <= maxSum) {
        int cnt_optleft = getSumOptLeft(2 * (maxSum - sum_in_path));
        res = cnt_in_path + cnt_optleft;
    }

    for (int i = 0; i < numNode; ++i) {
        int u(sorted_diff[i].node);
        if(inPath[u]) {
            sum_in_path += dist_h[u] - dist_l[u];
            ++cnt_in_path;
        } else {
            int pos = upper_bound(sorted_val.begin(), sorted_val.end(), pli(2 * dist_l[u], u)) - sorted_val.begin();
            modify(fen_dist, pos, -2 * dist_l[u]);
            modify(fen_l, pos, -1);

            pos = upper_bound(sorted_val.begin(), sorted_val.end(), pli(dist_h[u], -2 * u - 2)) - sorted_val.begin();
            modify(fen_dist, pos, dist_h[u]);
            modify(fen_h, pos, +1);

            modify(fen_dist, pos + 1, dist_h[u]);
            modify(fen_h, pos + 1, +1);
        }

        if(sum_in_path > maxSum)
            break;

        int cnt_optleft = getSumOptLeft(2 * (maxSum - sum_in_path));
        res = max(res, cnt_in_path + cnt_optleft);
    }

    return res;
}

int calcWithoutCommon(void) {
    vector<ll> sorted_dist;
    for (int i = 0; i < numNode; ++i) {
        sorted_dist.push_back(dist_l[i]);
        sorted_dist.push_back(dist_h[i]);
    }

    sort(sorted_dist.begin(), sorted_dist.end());

    ll tot_sum(0);
    int res(0);
    for (int i = 0; i < sz(sorted_dist); ++i) {
        ll dis(sorted_dist[i]);

        tot_sum += dis;
        if(tot_sum <= maxSum)
            res = i + 1;
    }

    return res;
}

int max_score(int _N, int _X, int _Y, ll _K, vector<int> _U, vector<int> _V, vector<int> _W) {
    numNode = _N, x_node = _X, y_node = _Y, maxSum = _K;
    for (int i = 0; i < numNode; ++i) {
        adj[i].clear();
        inPath[i] = 0;
    }

    for (int i = 0; i + 1 < numNode; ++i) {
        int u(_U[i]), v(_V[i]), w(_W[i]);
        adj[u].push_back(ii(v, w));
        adj[v].push_back(ii(u, w));
    }

    dfs_init(x_node, -1, 0, dist_l, pa);
    dfs_init(y_node, -1, 0, dist_h, pa);
    for (int u = x_node; u != -1; u = pa[u])
        inPath[u] = 1;

    for (int i = 0; i < numNode; ++i) {
        ll min_dist = min(dist_l[i], dist_h[i]);
        ll max_dist = max(dist_l[i], dist_h[i]);
        dist_l[i] = min_dist, dist_h[i] = max_dist;
    }

    return max(calcWithCommon(), calcWithoutCommon());
}

#ifdef Nhoksocqt1

int main(void) {
    ios_base::sync_with_stdio(0), cin.tie(0), cout.tie(0);

    #define TASK "closing"
    if(fopen(TASK".inp", "r")) {
        freopen(TASK".inp", "r", stdin);
        freopen(TASK".out", "w", stdout);
    }

    vector<int> _U, _V, _W;
    int _N, _X, _Y, _K;
    cin >> _N >> _X >> _Y >> _K;

    _U.resize(_N), _V.resize(_N), _W.resize(_N);
    for (int i = 0; i + 1 < _N; ++i)
        cin >> _U[i] >> _V[i] >> _W[i];

    int x = max_score(_N, _X, _Y, _K, _U, _V, _W);
    cout << "ANSWER: " << x << '\n';

    return 0;
}

#endif // Nhoksocqt1

• Subtask 0: 0 / 0
• Subtask 1: 8 / 8
• Subtask 2: 9 / 9
• Subtask 3: 12 / 12
• Subtask 4: 14 / 14
• Subtask 5: 9 / 9
• Subtask 6: 11 / 11
• Subtask 7: 10 / 10
• Subtask 8: 10 / 10
• Subtask 9: 17 / 17