Submission #788531

#	Time	Username	Problem	Language	Result	Execution time	Memory
788531		79brue	Paths (RMI21_paths)	C++17	100 / 100	481 ms	64836 KiB

Main

#include <bits/stdc++.h>

using namespace std;

typedef long long ll;

int cnt[800002]; ll sum[800002];

inline void update(int i, int l, int r, int x, int v1, ll v2){
    if(l==r){
        cnt[i] += v1, sum[i] += v2;
        return;
    }
    int m = (l+r)>>1;
    if(x<=m) update(i*2, l, m, x, v1, v2);
    else update(i*2+1, m+1, r, x, v1, v2);
    cnt[i] = cnt[i*2] + cnt[i*2+1], sum[i] = sum[i*2] + sum[i*2+1];
}

inline ll query(int i, int l, int r, int x){
    if(l==r) return cnt[i] ? (sum[i] / cnt[i]) * x : 0;
    int m = (l+r)>>1;
    if(cnt[i*2+1] >= x) return query(i*2+1, m+1, r, x);
    else return query(i*2, l, m, x-cnt[i*2+1]) + sum[i*2+1];
}

struct Edge{
    int s, e; ll v; int idx;
    Edge(){}
    Edge(int s, int e, ll v): s(s), e(e), v(v){}
    bool operator<(const Edge &r)const{
        return idx<r.idx;
    }
};

int n, k;
Edge arr[200002];
vector<Edge> link[100002];
int in[100002], out[100002], idx[100002], inCnt;
int par[100002];

void dfs_in(int x, int p=-1){
    in[x] = ++inCnt;
    idx[inCnt] = x;
    par[x] = p;
    for(auto y: link[x]){
        if(y.e == p) continue;
        dfs_in(y.e, x);
    }
    out[x] = inCnt;
}

vector<Edge> linkSet[100002];
ll MX[200002], where[200002]; /// 이 간선 방향으로 갔을 때 최대가 몇인가
vector<pair<ll, ll> > options[100002]; /// 이 정점에서 나갈 수 있는 모든 옵션

pair<ll, int> dfs_getValues(int x, int p=-1){
//    printf("Get value %d %d\n", x, p);
    if(!linkSet[x].empty()){
        vector<Edge> tlst;
        for(Edge y: linkSet[x]){
            if((y.idx ^ p) == 1){
                tlst.push_back(y);
                continue;
            }
            options[x].push_back(make_pair(MX[y.idx] = (dfs_getValues(y.e, y.idx).first + y.v), y.idx));
        }
        linkSet[x].swap(tlst);
        sort(options[x].rbegin(), options[x].rend());
    }
    if(options[x].empty()) return make_pair(0, -1);
    else if((options[x][0].second ^ p) != 1) return options[x][0];
    else if((int)options[x].size() == 1) return make_pair(0, -1);
    else return options[x][1];
}

ll numbers[200002];
vector<ll> inQuery[100002], outQuery[100002];
ll ans[100002];

void putQuery(int s, int e, ll p, int mode){
    if(in[s] < in[e]){
        if(mode == 1){
            inQuery[1].push_back(p), outQuery[in[e]].push_back(p);
            if(out[e]+1 <= n+1) inQuery[out[e]+1].push_back(p), outQuery[n+1].push_back(p);
        }
        else{
            outQuery[1].push_back(p), inQuery[in[e]].push_back(p);
            if(out[e]+1 <= n+1) outQuery[out[e]+1].push_back(p), inQuery[n+1].push_back(p);
        }
    }
    else{
        if(mode == 1) inQuery[in[s]].push_back(p), outQuery[out[s]+1].push_back(p);
        else          outQuery[in[s]].push_back(p), inQuery[out[s]+1].push_back(p);
    }
}

int main(){
    scanf("%d %d", &n, &k);
    for(int i=1; i<n; i++){
        scanf("%d %d %lld", &arr[i*2-2].s, &arr[i*2-2].e, &arr[i*2-2].v);
        arr[i*2-1].s = arr[i*2-2].e, arr[i*2-1].e = arr[i*2-2].s, arr[i*2-1].v = arr[i*2-2].v;
        arr[i*2-2].idx = i*2-2, arr[i*2-1].idx = i*2-1;
        link[arr[i*2-2].s].push_back(arr[i*2-2]);
        link[arr[i*2-1].s].push_back(arr[i*2-1]);
    }
    for(int i=1; i<=n; i++) linkSet[i] = link[i];

    dfs_in(1);
    for(int i=0; i<(n-1)*2; i++){
        if(MX[i]) continue;
        pair<ll, int> p = dfs_getValues(arr[i].e, i);
        MX[i] = p.first + arr[i].v;
        where[i] = p.second;
    }
    for(int i=1; i<=n; i++) options[i].clear();
    for(int i=0; i<(n-1)*2; i++){
        options[arr[i].s].push_back(make_pair(MX[i], i));
    }
    for(int i=1; i<=n; i++) sort(options[i].rbegin(), options[i].rend());

    for(int i=1; i<=n; i++){
        /// 수는 최대 2N개
        for(int j=0; j<(int)options[i].size(); j++){
            int p = options[i][j].second; ll v = options[i][j].first;
            int s = arr[p].s, e = arr[p].e;
            putQuery(s, e, v, 1);
            if(where[p] != -1 && v-arr[p].v) putQuery(s, e, v-arr[p].v, -1);
        }
    }

    numbers[0] = 0;
    for(int i=0; i<(n-1)*2; i++) numbers[i+1] = MX[i];
    sort(numbers, numbers+n*2-1);
    int L = unique(numbers, numbers+n*2-1) - numbers;
    for(int i=1; i<=n; i++) for(ll &p: inQuery[i]) p = lower_bound(numbers, numbers+L, p) - numbers;
    for(int i=1; i<=n; i++) for(ll &p: outQuery[i]) p = lower_bound(numbers, numbers+L, p) - numbers;

    for(int i=1; i<=n; i++){
        for(ll p: inQuery[i]){
//            printf("In update %d %lld\n", i, numbers[p]);
            update(1, 0, L-1, p, 1, numbers[p]);
        }
        for(ll p: outQuery[i]){
//            printf("Out update %d %lld\n", i, numbers[p]);
            update(1, 0, L-1, p, -1, -numbers[p]);
        }
        ans[idx[i]] = query(1, 0, L-1, k);
    }

    for(int i=1; i<=n; i++){
        printf("%lld\n", ans[i]);
    }
}

Compilation message (stderr)

Main.cpp: In function 'int main()':
Main.cpp:99:10: warning: ignoring return value of 'int scanf(const char*, ...)' declared with attribute 'warn_unused_result' [-Wunused-result]
   99 |     scanf("%d %d", &n, &k);
      |     ~~~~~^~~~~~~~~~~~~~~~~
Main.cpp:101:14: warning: ignoring return value of 'int scanf(const char*, ...)' declared with attribute 'warn_unused_result' [-Wunused-result]
  101 |         scanf("%d %d %lld", &arr[i*2-2].s, &arr[i*2-2].e, &arr[i*2-2].v);
      |         ~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

• Subtask 1: 8 / 8
• Subtask 2: 11 / 11
• Subtask 3: 17 / 17
• Subtask 4: 20 / 20
• Subtask 5: 12 / 12
• Subtask 6: 32 / 32