oj.uz

Submission #381957

# Submission time Handle Problem Language Result Execution time Memory
381957 2021-03-26T07:43:21 Z Vimmer Janjetina (COCI21_janjetina) C++14
110 / 110
 1261 ms 21096 KB 
#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;

        rec(u, max(it.S, 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);
    }

    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);
                }

            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;
    }

    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);
                    }

                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);
        }
    }

    assert(sum % 2 == 0);

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

# Verdict Execution time Memory Grader output
1 Correct 2 ms 2668 KB Output is correct
2 Correct 2 ms 2668 KB Output is correct
3 Correct 3 ms 2668 KB Output is correct
4 Correct 7 ms 2924 KB Output is correct
5 Correct 5 ms 2924 KB Output is correct
6 Correct 6 ms 2924 KB Output is correct
7 Correct 6 ms 2924 KB Output is correct
8 Correct 7 ms 2924 KB Output is correct
9 Correct 4 ms 2796 KB Output is correct
10 Correct 4 ms 2796 KB Output is correct
11 Correct 4 ms 2796 KB Output is correct
12 Correct 6 ms 2796 KB Output is correct
13 Correct 6 ms 2796 KB Output is correct
14 Correct 7 ms 2796 KB Output is correct
15 Correct 5 ms 2796 KB Output is correct
16 Correct 5 ms 2796 KB Output is correct
17 Correct 6 ms 2796 KB Output is correct
18 Correct 6 ms 2812 KB Output is correct
19 Correct 6 ms 2796 KB Output is correct
20 Correct 6 ms 2796 KB Output is correct
21 Correct 6 ms 2796 KB Output is correct
22 Correct 6 ms 2796 KB Output is correct

Subtask #2
35.0 / 35.0

# Verdict Execution time Memory Grader output
1 Correct 3 ms 2668 KB Output is correct
2 Correct 2 ms 2668 KB Output is correct
3 Correct 3 ms 2668 KB Output is correct
4 Correct 7 ms 2924 KB Output is correct
5 Correct 71 ms 4588 KB Output is correct
6 Correct 538 ms 12012 KB Output is correct
7 Correct 797 ms 20832 KB Output is correct
8 Correct 894 ms 20936 KB Output is correct
9 Correct 976 ms 20956 KB Output is correct
10 Correct 1064 ms 20832 KB Output is correct
11 Correct 863 ms 20776 KB Output is correct
12 Correct 1028 ms 20836 KB Output is correct
13 Correct 930 ms 20944 KB Output is correct
14 Correct 1112 ms 20820 KB Output is correct
15 Correct 1072 ms 21096 KB Output is correct
16 Correct 1233 ms 21036 KB Output is correct
17 Correct 1228 ms 21024 KB Output is correct
18 Correct 1220 ms 20980 KB Output is correct

Subtask #3
60.0 / 60.0

# Verdict Execution time Memory Grader output
1 Correct 2 ms 2668 KB Output is correct
2 Correct 2 ms 2668 KB Output is correct
3 Correct 3 ms 2668 KB Output is correct
4 Correct 7 ms 2924 KB Output is correct
5 Correct 5 ms 2924 KB Output is correct
6 Correct 6 ms 2924 KB Output is correct
7 Correct 6 ms 2924 KB Output is correct
8 Correct 7 ms 2924 KB Output is correct
9 Correct 4 ms 2796 KB Output is correct
10 Correct 4 ms 2796 KB Output is correct
11 Correct 4 ms 2796 KB Output is correct
12 Correct 6 ms 2796 KB Output is correct
13 Correct 6 ms 2796 KB Output is correct
14 Correct 7 ms 2796 KB Output is correct
15 Correct 5 ms 2796 KB Output is correct
16 Correct 5 ms 2796 KB Output is correct
17 Correct 6 ms 2796 KB Output is correct
18 Correct 6 ms 2812 KB Output is correct
19 Correct 6 ms 2796 KB Output is correct
20 Correct 6 ms 2796 KB Output is correct
21 Correct 6 ms 2796 KB Output is correct
22 Correct 6 ms 2796 KB Output is correct
23 Correct 3 ms 2668 KB Output is correct
24 Correct 2 ms 2668 KB Output is correct
25 Correct 3 ms 2668 KB Output is correct
26 Correct 7 ms 2924 KB Output is correct
27 Correct 71 ms 4588 KB Output is correct
28 Correct 538 ms 12012 KB Output is correct
29 Correct 797 ms 20832 KB Output is correct
30 Correct 894 ms 20936 KB Output is correct
31 Correct 976 ms 20956 KB Output is correct
32 Correct 1064 ms 20832 KB Output is correct
33 Correct 863 ms 20776 KB Output is correct
34 Correct 1028 ms 20836 KB Output is correct
35 Correct 930 ms 20944 KB Output is correct
36 Correct 1112 ms 20820 KB Output is correct
37 Correct 1072 ms 21096 KB Output is correct
38 Correct 1233 ms 21036 KB Output is correct
39 Correct 1228 ms 21024 KB Output is correct
40 Correct 1220 ms 20980 KB Output is correct
41 Correct 2 ms 2668 KB Output is correct
42 Correct 791 ms 21072 KB Output is correct
43 Correct 893 ms 20832 KB Output is correct
44 Correct 891 ms 21088 KB Output is correct
45 Correct 1057 ms 20884 KB Output is correct
46 Correct 784 ms 20832 KB Output is correct
47 Correct 922 ms 20920 KB Output is correct
48 Correct 897 ms 20840 KB Output is correct
49 Correct 1041 ms 20996 KB Output is correct
50 Correct 1060 ms 20832 KB Output is correct
51 Correct 1194 ms 20960 KB Output is correct
52 Correct 397 ms 14048 KB Output is correct
53 Correct 430 ms 14048 KB Output is correct
54 Correct 298 ms 14140 KB Output is correct
55 Correct 430 ms 14176 KB Output is correct
56 Correct 437 ms 14048 KB Output is correct
57 Correct 1061 ms 13540 KB Output is correct
58 Correct 1078 ms 13668 KB Output is correct
59 Correct 856 ms 13540 KB Output is correct
60 Correct 1234 ms 13540 KB Output is correct
61 Correct 1261 ms 13796 KB Output is correct
62 Correct 835 ms 13796 KB Output is correct
63 Correct 981 ms 13672 KB Output is correct
64 Correct 1097 ms 13536 KB Output is correct
65 Correct 28 ms 3308 KB Output is correct
66 Correct 2 ms 2668 KB Output is correct