Submission #396234

#	Time	Username	Problem	Language	Result	Execution time	Memory
396234		MrRobot_28	Janjetina (COCI21_janjetina)	C++17	110 / 110	1022 ms	58644 KiB

Main

#include<bits/stdc++.h>
using namespace std;
#define X first
#define Y second
#define sz(a) (int)a.size()
#define ll long long
const int N = 1e5;
const int T = 20;
vector <pair <int, int> > g[N];
int n, k;
int used[N];
int _sz[N];
vector <int> h_p[N];
vector <int> Trees[N];
ll ans = 0;
void update(int ind, int v, int l, int r, int x, int c)
{
    Trees[ind][v] += c;
    if(l == r)
    {
        return;
    }
    if(x <= (r + l) / 2)
    {
        update(ind, v * 2, l, (r + l) / 2, x, c);
    }
    else
    {
        update(ind, v * 2 + 1, (r + l) / 2 + 1, r, x, c);
    }
}
int get_ans(int ind, int v, int l, int r, int al, int ar)
{
    if(al > ar)
    {
        return 0;
    }
    if(l >= al && r <= ar)
    {
        return Trees[ind][v];
    }
    if(l <= ar && r >= al)
    {
        return get_ans(ind, v * 2, l, (r + l) / 2, al, ar) + get_ans(ind, v * 2 + 1, (r + l) / 2 + 1, r, al, ar);
    }
    return 0;
}

int coun[N];
void dfs(int v, int p = -1)
{
    _sz[v] = 1;
    for(auto to : g[v])
    {
        if(!used[to.X] && to.X != p)
        {
            dfs(to.X, v);
            _sz[v] += _sz[to.X];
        }
    }
}
int centroid(int v, int sz1, int p = -1)
{
    for(auto to : g[v])
    {
        if(!used[to.X] && to.X != p && _sz[to.X] * 2 > sz1)
        {
            return centroid(to.X, sz1, v);
        }
    }
    return v;
}
int par[N];
int H[N];
vector <pair <int, int> > vec;
vector <int> h_t[N];
void go_to(int v, int maxd, int h, int p = -1)
{
    if(maxd >= h)
    {
        ans++;
    }
    h_t[par[v]].push_back(h);
    H[v] = h;
    coun[par[v]]++;
    vec.push_back({maxd, v});
    for(auto to : g[v])
    {
        if(used[to.X] || to.X == p)
        {
            continue;
        }
        par[to.X] = par[v];
        go_to(to.X, max(maxd, to.Y), h + 1, v);
    }
}

void build(int v)
{
    dfs(v);
    v = centroid(v, _sz[v]);
    used[v] = 1;
    for(auto to : g[v])
    {
        h_t[to.X].clear();
        coun[to.X] = 0;
        Trees[to.X + 1].clear();
        if(!used[to.X])
        {
            par[to.X] = to.X;
            go_to(to.X, to.Y, 1);
            sort(h_t[to.X].begin(), h_t[to.X].end());
            Trees[to.X + 1].resize(4 * coun[to.X]);
        }
    }
    sort(vec.begin(), vec.end());
    for(auto p : vec)
    {
        int t_l1 = p.X - H[p.Y];
        int t_l2 = upper_bound(h_t[par[p.Y]].begin(), h_t[par[p.Y]].end(), t_l1) - h_t[par[p.Y]].begin();
        t_l2--;
        ans += get_ans(0, 1, 0, N - 1, 0, t_l1) - get_ans(par[p.Y] + 1, 1, 0, coun[par[p.Y]] - 1, 0, t_l2);
        update(0, 1, 0, N - 1, H[p.Y], 1);
        int d1 = lower_bound(h_t[par[p.Y]].begin(), h_t[par[p.Y]].end(), H[p.Y]) - h_t[par[p.Y]].begin();
        update(par[p.Y] + 1, 1, 0, coun[par[p.Y]] - 1, d1, 1);
    }
 //   cout << v << " " << ans << "\n";
    for(auto p : vec)
    {
        int t_l1 = p.X - H[p.Y];
        int t_l2 = upper_bound(h_t[par[p.Y]].begin(), h_t[par[p.Y]].end(), t_l1) - h_t[par[p.Y]].begin();
        t_l2--;
        update(0, 1, 0, N - 1, H[p.Y], -1);
        int d1 = lower_bound(h_t[par[p.Y]].begin(), h_t[par[p.Y]].end(), H[p.Y]) - h_t[par[p.Y]].begin();
        update(par[p.Y] + 1, 1, 0, coun[par[p.Y]] - 1, d1, -1);
    }
    vec.clear();
    for(auto to : g[v])
    {
        if(!used[to.X])
        {
            build(to.X);
        }
    }
}
signed main()
{
  //  ifstream cin("input1.txt.4c");
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);
    cin >> n >> k;
    for(int i = 0; i < n - 1; i++)
    {
        int x, y, w;
        cin >> x >> y >> w;
        x--;
        y--;
        w -= k;
        g[x].push_back({y, w});
        g[y].push_back({x, w});
    }
    Trees[0].resize(4 * N);
    build(0);
    cout << ans * 2;
    return 0;
}

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