Submission #1064393

#	Time	Username	Problem	Language	Result	Execution time	Memory
1064393		rembocoder	Janjetina (COCI21_janjetina)	C++17	110 / 110	736 ms	28208 KiB

Main

#include <bits/stdc++.h>
#include<ext/pb_ds/assoc_container.hpp>

using namespace std;
using namespace __gnu_pbds;

#define int int64_t

typedef tree<pair<int, int>, null_type, less<pair<int, int>>, rb_tree_tag, tree_order_statistics_node_update> ordered_set;

const int inf = 2e9;
const int mod = 1e9 + 7;

int ans = 0;
int k;
vector<vector<pair<int, int>>> g;
vector<vector<int>> centers;
vector<bool> del;
vector<int> sz;

void find_sz(int v, int p) {
    sz[v] = 1;
    for (auto to_w: g[v]) {
        int to = to_w.first, w = to_w.second;
        if (to == p || del[to]) {
            continue;
        }
        find_sz(to, v);
        sz[v] += sz[to];
    }
}

int find_center(int v, int p, int total_sz) {
    for (auto to_w: g[v]) {
        int to = to_w.first, w = to_w.second;
        if (to == p || del[to]) {
            continue;
        }
        if (sz[to] * 2 > total_sz) {
            return find_center(to, v, total_sz);
        }
    }
    return v;
}

void fill_list(int v, int p, int cur_max, int cur_dist, vector<pair<int, int>>& lst) {
    lst.push_back({cur_max, cur_dist});
    for (auto to_w: g[v]) {
        int to = to_w.first, w = to_w.second;
        if (to == p || del[to]) {
            continue;
        }
        fill_list(to, v, max(cur_max, w), cur_dist + 1, lst);
    }
}

void handle_center(int v) {
    del[v] = true;
    vector<vector<pair<int, int>>> lists; // {max_on_path, dist}
    vector<pair<int, int>> common_list;
    common_list.push_back({-inf, 0});
    for (auto to_w: g[v]) {
        int to = to_w.first, w = to_w.second;
        lists.push_back(vector<pair<int, int>>(0));
        if (!del[to]) {
            fill_list(to, -1, w, 1, lists.back());
        }
        sort(lists.back().begin(), lists.back().end());
        common_list.insert(common_list.end(), lists.back().begin(), lists.back().end());
    }
    sort(common_list.begin(), common_list.end());
    auto handle_list = [&](vector<pair<int, int>>& lst) {
        int counter = 0;
        int cur_ans = 0;
        ordered_set s;
        for (auto mx_dist: lst) {
            int mx = mx_dist.first, dist = mx_dist.second;
            cur_ans += s.order_of_key({mx - k - dist + 1, -1});
            s.insert({dist, counter++});
        }
        return cur_ans;
    };
    ans += handle_list(common_list);
    for (auto lst: lists) {
        ans -= handle_list(lst);
    }
}

void centroid_decomposition(int v) {
    find_sz(v, -1);
    v = find_center(v, -1, sz[v]);
    handle_center(v);
    del[v] = true;
    for (auto to_w: g[v]) {
        int to = to_w.first, w = to_w.second;
        if (!del[to]) {
            centroid_decomposition(to);
        }
    }
}

int32_t main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0); cout.tie(0);
    int n;
    cin >> n >> k;
    g.resize(n);
    for (int i = 0; i < n - 1; i++) {
        int a, b, c;
        cin >> a >> b >> c;
        a--; b--;
        g[a].push_back({b, c});
        g[b].push_back({a, c});
    }
    del.resize(n);
    sz.resize(n);
    centers.resize(n);
    centroid_decomposition(0);
    cout << ans * 2 << '\n';
}

Compilation message (stderr)

Main.cpp: In function 'void find_sz(int64_t, int64_t)':
Main.cpp:24:30: warning: unused variable 'w' [-Wunused-variable]
   24 |         int to = to_w.first, w = to_w.second;
      |                              ^
Main.cpp: In function 'int64_t find_center(int64_t, int64_t, int64_t)':
Main.cpp:35:30: warning: unused variable 'w' [-Wunused-variable]
   35 |         int to = to_w.first, w = to_w.second;
      |                              ^
Main.cpp: In function 'void centroid_decomposition(int64_t)':
Main.cpp:95:30: warning: unused variable 'w' [-Wunused-variable]
   95 |         int to = to_w.first, w = to_w.second;
      |                              ^

• Subtask 1: 15 / 15
• Subtask 2: 35 / 35
• Subtask 3: 60 / 60