oj.uz

Submission #541693

# Submission time Handle Problem Language Result Execution time Memory
541693 2022-03-24T07:57:53 Z AlperenT Janjetina (COCI21_janjetina) C++17
0 / 110
 2 ms 2684 KB 
#include <bits/stdc++.h>

using namespace std;

const int N = 1e5 + 5;

int n, k, a, b, w, treesz[N];

bool isblocked[N];

struct Edge{
    int b, w;
};

vector<Edge> tree[N];

struct Item{
    int w, d;

    bool operator < (const Item &sc) const{
        if(w == sc.w) return d < sc.d;
        else return w < sc.w;
    }
};

struct Fenwick{
    int tree[N];

    stack<pair<int, int>> hist;

    int getsum(int r){
        int sum = 0;
        for(int i = r; i > 0; i -= i & -i) sum += tree[i];
        return sum;
    }

    int query(int l, int r){
        if(l > r) return 0;
        return getsum(r) - getsum(l - 1);
    }

    void update(int pos, int val, bool addhist = true){
        if(addhist) hist.push({pos, val});
        for(int i = pos; i < N; i += i & -i) tree[i] += val;
    }

    void clear(){
        while(!hist.empty()){
            update(hist.top().first, -hist.top().second, false);
            hist.pop();
        }
    }
};

Fenwick fw;

int calcsz(int v, int p){
    treesz[v] = 1;

    for(auto e : tree[v]){
        if(e.b != p && !isblocked[e.b]) treesz[v] += calcsz(e.b, v);
    }

    return treesz[v];
}

int findcent(int v, int p, int sz){
    for(auto e : tree[v]){
        if(e.b != p && !isblocked[e.b] && treesz[e.b] >= sz / 2){
            return findcent(e.b, v, sz);
        }
    }

    return v;
}

void dfs(int v, int p, int d, int w, vector<Item> &vec){
    vec.push_back({w, d});

    for(auto e : tree[v]){
        if(e.b != p && !isblocked[e.b]){
            dfs(e.b, v, d + 1, max(w, e.w), vec);
        }
    }
}

long long centdecomp(int v){
    int sz = calcsz(v, 0);

    int cent = findcent(v, 0, sz);

    long long ans = 0;

    vector<Item> total, cur;

    for(auto e : tree[cent]){
        if(!isblocked[e.b]){
            cur.clear();

            dfs(e.b, cent, 1, e.w, cur);

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

            for(auto itm : cur){
                if(itm.w - itm.d >= k) ans += 1;

                ans -= fw.query(0, itm.w - itm.d - k);

                fw.update(itm.d, 1);

                total.push_back(itm);
            }

            fw.clear();
        }
    }

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

    for(auto itm : total){
        ans += fw.query(0, itm.w - itm.d - k);

        fw.update(itm.d, 1);
    }

    fw.clear();

    isblocked[cent] = true;

    for(auto e : tree[v]){
        if(!isblocked[e.b]){
            ans += centdecomp(e.b);
        }
    }

    return ans;
}

int main(){
    ios_base::sync_with_stdio(false);cin.tie(NULL);

    cin >> n >> k;

    for(int i = 1; i <= n - 1; i++){
        cin >> a >> b >> w;

        tree[a].push_back({b, w});
        tree[b].push_back({a, w});
    }

    cout << centdecomp(1) * 2;
}

Subtask #1
0 / 15.0

# Verdict Execution time Memory Grader output
1 Correct 2 ms 2644 KB Output is correct
2 Correct 1 ms 2644 KB Output is correct
3 Incorrect 2 ms 2672 KB Output isn't correct
4 Halted 0 ms 0 KB -

Subtask #2
0 / 35.0

# Verdict Execution time Memory Grader output
1 Correct 2 ms 2644 KB Output is correct
2 Correct 2 ms 2644 KB Output is correct
3 Incorrect 2 ms 2684 KB Output isn't correct
4 Halted 0 ms 0 KB -

Subtask #3
0 / 60.0

# Verdict Execution time Memory Grader output
1 Correct 2 ms 2644 KB Output is correct
2 Correct 1 ms 2644 KB Output is correct
3 Incorrect 2 ms 2672 KB Output isn't correct
4 Halted 0 ms 0 KB -