Submission #381922

#	Time	Username	Problem	Language	Result	Execution time	Memory
381922		Vimmer	Janjetina (COCI21_janjetina)	C++14	0 / 110	3 ms	2668 KiB

Main

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

#define N 100500
#define PB push_back
#define endl '\n'
#define pri(x) cout << x << endl
#define _ << " " <<
#define sz(x) int(x.size())
#define F first
#define S second
#define all(x) x.begin(), x.end()

using namespace std;
using namespace __gnu_pbds;

typedef long long ll;

typedef tree <int, null_type, less_equal <int>, rb_tree_tag, tree_order_statistics_node_update> ord_set;

ll ans = 0;

ord_set se;

vector <pair <int, int> > vr;

vector <pair <int, int> > g[N];

vector <int> vt;

int siz[N], n, k, a[N], h[N];

bool mk[N];

bool cmp(int x, int y)
{
    return a[x] < a[y];
}

void dfs(int v, int p)
{
    siz[v] = 1;

    for (auto it : g[v])
    {
        int u = it.F;

        if (u == p) continue;

        if (mk[u]) continue;

        dfs(u, v);

        siz[v] += siz[u];
    }
}

int fnd_root(int v)
{
    for (auto it : g[v])
      if (!mk[it.F] && siz[v] / 2 < siz[it.F])
      {
          siz[v] -= siz[it.F];

          siz[it.F] += siz[v];

          return fnd_root(it.F);
      }

    return v;
}

void rec(int v, int mx, int ht, int p)
{
    h[v] = ht;

    a[v] = mx;

    vr.PB({mx, ht + k});

    vt.PB(v);

    for (auto it : g[v])
    {
        int u = it.F;

        if (mk[u] || u == p) continue;

        mx = max(mx, it.S);

        rec(u, mx, ht + 1, v);
    }
}

void dfser(int v, int mx, int ht, int p)
{
    if (mx - ht >= k)
        ans += 2;

    for (auto it : g[v])
    {
        if (mk[it.F] || it.F == p) continue;

        dfser(it.F, max(mx, it.S), ht + 1, v);
    }
}

void calc(int v)
{
    mk[v] = 1;

    vr.clear(); vt.clear();

    se.clear();

    for (auto it : g[v])
    {
        if (mk[it.F]) continue;

        rec(it.F, it.S, 1, -1);
    }

    /// for root

    sort(all(vt), cmp);

    sort(all(vr));

    int j = 0;

    ll sum = 0;

    for (auto v : vt)
    {
        while (j < sz(vr) && vr[j].F < a[v])
        {
            se.insert(vr[j].S);

            j++;
        }

        int vl = a[v] - h[v];

        int kl = se.order_of_key(vl + 1);

        sum += kl + kl;
    }

   // pri("DOING:" _  v _ "SUM:" _ sum _ "ANS:" _ ans);

    for (auto it : g[v])
    {
        int u = it.F;

        if (mk[u]) continue;

        vr.clear(); se.clear();

        vt.clear();

        rec(u, it.S, 1, -1);

        sort(all(vt), cmp);

        sort(all(vr));

        int j = 0;

        for (auto v : vt)
        {
            while (j < sz(vr) && vr[j].F <= a[v])
            {
                se.insert(vr[j].S);

                j++;
            }

            int vl = a[v] - h[v];

            int kl = se.order_of_key(vl + 1);

            sum -= kl - kl;

            //cout << "DEL: " << v << " " << kl << endl;
        }
    }

    //assert(sum % 2 == 0);

    //pri("DOING:" _  v _ "SUM:" _ sum _ "ANS:" _ ans);
    ans += sum;

    for (auto it : g[v])
    {
        int u = it.F;

        if (mk[u]) continue;

        dfser(u, it.S, 1, -1);
    }

    for (auto it : g[v])
    {
        if (mk[it.F]) continue;

        int root = fnd_root(it.F);

        dfs(root, -1);

        calc(root);
    }
}

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

//    freopen("1.in", "r", stdin);

    cin >> n >> k;

    for (int i = 1; i < n; i++)
    {
        int x, y, z;

        cin >> x >> y >> z;

        g[x].PB({y, z});

        g[y].PB({x, z});
    }

    dfs(1, -1);

    int root = fnd_root(1);

    dfs(root, -1);

    calc(root);

    pri(ans);
}

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