제출 #381944

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
381944		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, st;

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(); st.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, lst = -1, u = 0;

    ll sum = 0;

//    cout << "DOING: " << v << ": ";

    for (auto v : vt)
    {
        if (lst != a[v])
            st.clear();

        while (u < sz(vr) && vr[u].F <= a[v])
        {
            if (vr[u].F == a[v])
                {
                    st.insert(vr[u].S);

//                    cout << "DOBR: " << vr[u].S << endl;
                }

            u++;
        }

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

            j++;
        }

        lst = a[v];

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

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

        sum += kl + kl;

        int db = st.order_of_key(vl + 1);
        sum += db;

//        cout << "AD: " << vl _ db << endl;
    }

//    cout << endl;

   // 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(); st.clear();

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

        sort(all(vt), cmp);

        sort(all(vr));

        int j = 0, lst = -1, p = 0;

        for (auto v : vt)
        {
            if (a[v] != lst)
                st.clear();

            while (p < sz(vr) && vr[p].F <= a[v])
            {
                if (vr[p].F == a[v])
                    {
                        st.insert(vr[p].S);

    //                    cout << "DOBR: " << vr[p].S << endl;
                    }

                p++;
            }

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

                j++;
            }

            lst = a[v];

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

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

            sum -= kl - kl;

            sum -= st.order_of_key(vl + 1);

            //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