#include <iostream>
#include <vector>
#include <assert.h>
using namespace std;

//#define int long long
typedef long long ll;
#define pii pair<ll, ll>
#define all(v) v.begin(), v.end()
#define ff first
#define ss second
const ll MAXN = 2e5 + 5;
const ll MAXM = 5e6 + 5;   // implicit segtree
const ll INF = 1e18 + 7;
const ll MAX = 1e12 + 5;

ll dist[MAXN][2];
ll a[MAXN], b[MAXN];   // 2 costs
ll ord[MAXN];
ll ans, ans2;  // no overlap, overlap
vector<pii> adj[MAXN];

ll have;

ll tot = 1;
ll val[MAXM], cnt[MAXM];
ll lc[MAXM], rc[MAXM];
vector<ll> costa, costb;

bool comp(ll i, ll j) {
    if(costb[i] == costb[j]) return costa[i] < costa[j];
    return costb[i] < costb[j];
}

void check() {
    if(tot >= MAXM) while(true) cout << "oof";
}

void pull(ll u) {
    val[u] = val[lc[u]] + val[rc[u]];
    cnt[u] = cnt[lc[u]] + cnt[rc[u]];
}

void ins(ll u, ll l, ll r, ll pos, ll freq) {  // segtree indexed by val so pos = val
    check();
    if(l == r) {
//        cout << "updating with " << pos << " " << freq << endl;
        val[u] += pos * freq;
        cnt[u] += freq;
        return;
    }

    ll mid = (l + r) / 2;

    if(pos <= mid) {
        if(lc[u] == 0) lc[u] = ++tot;
        ins(lc[u], l, mid, pos, freq);
    }
    else {
        if(rc[u] == 0) rc[u] = ++tot;
        ins(rc[u], mid + 1, r, pos, freq);
    }
    pull(u);
}

int walk(ll u, ll l, ll r) {
    if(u == 0) return 0;
//    cout << u << " " << l << " " << r << " " << val[u] << " " << cnt[u] << endl;
    if(l == r) {
//        cout << l << " " << cnt[u] << " " << have << " " << l << endl;
        ll sub = min(cnt[u], have / l);
        have -= l * sub;
//        cout << "returning " << sub << endl;
        return sub;
    }

    ll mid = (l + r) / 2;
    if(val[lc[u]] > have) return walk(lc[u], l, mid);

//    cout << "left val: " << val[lc[u]] << " node, lc: " << u << " " << lc[u] << endl;
    have -= val[lc[u]];
    if(have) return cnt[lc[u]] + walk(rc[u], mid + 1, r);
    return cnt[lc[u]];
}

void dfs(ll u, ll p, ll type) {
    for(auto [v, w] : adj[u]) {
        if(v == p) continue;
        dist[v][type] = dist[u][type] + w;

        dfs(v, u, type);
    }
}

int max_score(int n, int x, int y, ll k, vector<int> u, vector<int> v, vector<int> w) {
    for(int i = 0; i < n - 1; ++i) {
        adj[u[i]].push_back({v[i], w[i]});
        adj[v[i]].push_back({u[i], w[i]});
    }

    dfs(x, -1, 0); dfs(y, -1, 1);

    for(int i = 0; i < n; ++i) {
        a[i] = min(dist[i][0], dist[i][1]);
        b[i] = max(dist[i][0], dist[i][1]);
    }

    // seperate
    vector<ll> sortedDists;
    for(int i = 0; i < n; ++i) {
        sortedDists.push_back(dist[i][0]); sortedDists.push_back(dist[i][1]);
    }

    sort(all(sortedDists));

    ll tot = k;
    for(auto val : sortedDists) {
        if(tot >= val) {
            ++ans; tot -= val;
        }
        else break;
    }

    // overlap
    bool flag = false;
    ll distXY = dist[y][0];    // dist between x,y

//    cout << "dist between x,y: " << distXY << endl;

    int forced = 0;
    for(int i = 0; i < n; ++i) {
        ord[i] = i;
        if(dist[i][0] + dist[i][1] == distXY) { // on the path
//            cout << "node " << i << " is on path" << endl;
            k -= a[i]; ++forced;  // forced

//            cout << "i: " << i << endl;
            costa.push_back(b[i] - a[i]);
            costb.push_back(INF);
//            cout << b[i] - a[i] << endl;
            ins(1, 0, MAX, b[i] - a[i], 1);
            if(k < 0) flag = true;
//            cout << "need " << b[i] - a[i] << " to turn to 2" << endl;
        }
        else {
            costa.push_back(a[i]);
            costb.push_back(b[i]);
            ins(1, 0, MAX, a[i], 1);
        }
//            cout << "need " << a[i] << " for extra" << endl;
    }

    sort(ord, ord + n, comp);

//    cout << "ord: ";
//    for(int i = 0; i < n; ++i) cout << ord[i] << " "; cout << endl;
//    cout << "costa: ";
//    for(int i = 0; i < n; ++i) cout << costa[ord[i]] << " "; cout << endl;
//    cout << "costb: ";
//    for(int i = 0; i < n; ++i) cout << costb[ord[i]] << " "; cout << endl;
//    cout << endl;

    ll sum = 0;

    have = k;
    ans2 = forced + walk(1, 0, MAX);
//    cout << "allmax: " << walk(1, 0, MAX) << endl;

//    cout << "k: " << k << endl;
    for(int i = 0; i < n; ++i) {
//        cout << "before tot: " << t[1].cnt << endl;
        ins(1, 0, MAX, costa[ord[i]], -1); ins(1, 0, MAX, costb[ord[i]] - costa[ord[i]], 1);

        sum += costa[ord[i]];
        have = k - sum;
//        cout << "have: " << have << endl;
        if(have < 0) break;
        ll maxv = walk(1, 0, MAX) + i + 1;
//        cout << "maxv, forced: " << maxv << " " << forced << endl;
        ans2 = max(ans2, forced + maxv);
//        cout << "ans2: " << ans2 << endl;
    }

//    cout << "total val of tree: " << val[1] << endl;
//
//    cout << "forced overlap: " << forced << endl;
//    cout << "total overlap ans: " << ans2 << endl;
    if(flag) ans2 = 0;

//    cout << "have: " << have << endl;
//    cout << ans << " " << ans2 << endl;
    return max(ans, ans2);
}

answer.code: In function ‘int max_score(int, int, int, ll, std::vector<int>, std::vector<int>, std::vector<int>)’:
answer.code:114:5: error: ‘sort’ was not declared in this scope; did you mean ‘short’?
  114 |     sort(all(sortedDists));
      |     ^~~~
      |     short