Submission #929069

#	Time	Username	Problem	Language	Result	Execution time	Memory
929069		a_l_i_r_e_z_a	Janjetina (COCI21_janjetina)	C++17	110 / 110	499 ms	15832 KiB

Main

// In the name of God
// Hope is last to die

#include <bits/stdc++.h>
using namespace std;

typedef long long ll;
typedef long double ld;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;

#define pb push_back
// #define int long long    
#define S second
#define F first
#define mp make_pair
#define smax(x, y) (x) = max((x), (y))
#define smin(x, y) (x) = min((x), (y))
#define all(x) (x).begin(), (x).end()
#define len(x) ((int)(x).size())

const int maxn = 1e5 + 5;
const int inf = 1e9 + 7;
int n = 5, k = 1, sz[maxn], fen[maxn];
ll ans;
vector<pii> adj[maxn], vecf;
bool mark[maxn];

void upd(int i, int x){
    for(; i <= n; i += i & -i) fen[i] += x;
}

int get(int i){
    int res = 0;
    for(; i; i -= i & -i) res += fen[i];
    return res;
}

ll calc(vector<pii> vec){
    vector<int> h;
    h.pb(-inf);
    sort(all(vec));
    for(auto u: vec) h.pb(u.S);
    sort(all(h));
    h.resize(unique(all(h)) - h.begin());
    for(int i = 0; i < vec.size(); i++){
        int x = lower_bound(all(h), vec[i].S) - h.begin();
        upd(x, +1);
    }
    ll res = 0;
    for(int i = vec.size() - 1; i >= 0; i--){
        int x = lower_bound(all(h), vec[i].S) - h.begin();
        upd(x, -1);
        int j = upper_bound(all(h), vec[i].F - k - vec[i].S) - h.begin();
        j--;
        res += get(j);
    }
    return res;
}

int get_sz(int v, int p = -1){
    sz[v] = 1;
    for(auto [u, w]: adj[v]){
        if(u != p && !mark[u]) sz[v] += get_sz(u, v);
    }
    return sz[v];
}

int find(int v, int p, int s){
    for(auto [u, w]: adj[v]) if(u != p && !mark[u] && 2 * sz[u] > s) return find(u, v, s);
    return v;
}

void dfs(int v, int p, int h, int mx){
    vecf.pb(mp(mx, h));
    if(mx - k - h >= 0) ans++;
    for(auto [u, w]: adj[v]){
        if(u != p && !mark[u]){
            dfs(u, v, h + 1, max(mx, w));
        }
    }
}

void cd(int v){
    v = find(v, -1, get_sz(v));
    vector<pii> all;
    for(auto [u, w]: adj[v]){
        if(mark[u]) continue;
        vecf.clear();
        dfs(u, v, 1, w);
        ans -= calc(vecf);
        for(auto u: vecf) all.pb(u);
    }
    ans += calc(all);
    mark[v] = 1;
    for(auto [u, w]: adj[v]) if(!mark[u]) cd(u);
}

int32_t main()
{
    ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0);

    cin >> n >> k;
    for(int i = 0; i < n - 1; i++){
        int u, v, w; cin >> u >> v >> w;
        adj[u].pb(mp(v, w));
        adj[v].pb(mp(u, w));
    }
    cd(1);
    cout << ans * 2 << '\n';

    return 0;
}

Compilation message (stderr)

Main.cpp: In function 'll calc(std::vector<std::pair<int, int> >)':
Main.cpp:46:22: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   46 |     for(int i = 0; i < vec.size(); i++){
      |                    ~~^~~~~~~~~~~~

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