Submission #1348676

#TimeUsernameProblemLanguageResultExecution timeMemory
1348676haru09Janjetina (COCI21_janjetina)C++20
110 / 110
235 ms22376 KiB
#include<bits/stdc++.h>
using namespace std;
#define ll long long
#define fi first
#define se second
#define task "code"

const int ar = 1e5 + 5;
const ll mod = 1e9 + 7;
const int N = 1e5;
int n, K;
vector<pair<int, int>> ad[ar];
ll ans = 0;
int sz[ar];
bool del[ar];
void dfs(int u, int p)
{
    sz[u] = 1;
    for (auto [v, w] : ad[u])
    {
        if (v == p || del[v]) continue;
        dfs(v, u);
        sz[u] += sz[v];
    }
}
int centroid(int u, int p, int n)
{
    for (auto [v, w] : ad[u])
    {
        if (v == p || del[v]) continue;
        if (sz[v] > n / 2) return centroid(v, u, n);
    }
    return u;
}
struct BIT
{
    int bit[ar];
    queue<int> qu;
    void update(int u, int v)
    {
        while (u <= N)
        {
            bit[u] += v;
            u += u & (-u);
        }
    }
    int get(int u)
    {
        int ans = 0;
        while (u > 0)
        {
            ans += bit[u];
            u -= u & (-u);
        }
        return ans;
    }
    int get(int l, int r)
    {
        if (l > r) return 0;
        return get(r) - get(l - 1);
    }
    void reset()
    {
        while (qu.size())
        {
            bit[qu.front()] = 0;
            qu.pop();
        }
    }
} alltree, subtree;
vector<int> allVal, cand[ar];
int d[ar];
int timedfs = 0, L[ar], R[ar], node[ar], par[ar];
int base[ar], mx[ar];
bool check[ar];
void dfs_subtree(int u, int p)
{
    L[u] = ++timedfs;
    node[timedfs] = u;
    for (auto [v, w] : ad[u])
    {
        if (v == p || del[v]) continue;
        d[v] = d[u] + 1;
        par[v] = u;
        alltree.update(d[v], 1);
        allVal.push_back(w);
        cand[w].push_back(v);
        dfs_subtree(v, u);
    }
    R[u] = timedfs;
}
void dfs_smalltree(int u, int p, int w)
{
    int r = min(mx[u], w - d[u] - K + base[u]);
    int l = base[u] + 1;
    ans += alltree.get(1, w - d[u] - K) - subtree.get(l, r) + (w - d[u] - K >= 0);
    alltree.update(d[u], -1);
    subtree.update(d[u] + base[u], -1);
    check[u] = true;
    for (auto [v, w1] : ad[u])
    {
        if (v == p || del[v] || check[v]) continue;
        dfs_smalltree(v, u, w);
    }
}
void calc(int u)
{
    timedfs = 0;
    d[u] = 0;
    dfs_subtree(u, 0);
    for (auto [v, w] : ad[u])
    {
        if (del[v]) continue;
        for (int i = L[v]; i <= R[v]; i++)
        {
            int k = node[i];
            base[k] = L[v] - 1;
            mx[k] = R[v];
            subtree.update(base[k] + d[k], 1);
        }
    }
    sort(allVal.begin(), allVal.end());
    allVal.erase(unique(allVal.begin(), allVal.end()), allVal.end());
    reverse(allVal.begin(), allVal.end());
    for (auto w : allVal)
    {
        for (int v : cand[w])
        {
            if (check[v]) continue;
            dfs_smalltree(v, par[v], w);
        }
        cand[w].clear();
    }
    allVal.clear();
    alltree.reset();
    subtree.reset();
    for (int i = L[u]; i <= R[u]; i++)
    {
        int v = node[i];
        check[v] = false;
        d[v] = base[v] = mx[v] = par[v] = 0;
    }
}
void solve(int u)
{
    dfs(u, 0);
    int num = sz[u];
    u = centroid(u, 0, num);
    calc(u);
    del[u] = true;
    for (auto [v, w] : ad[u])
    {
        if (del[v]) continue;
        solve(v);
    }
}
int main()
{
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);
    if (fopen(task".inp", "r"))
    {
        freopen(task".inp", "r", stdin);
        freopen(task".out", "w", stdout);
    }
    cin >> n >> K;
    for (int i = 1; i < n; i++)
    {
        int u, v, w;
        cin >> u >> v >> w;
        ad[u].push_back({v, w});
        ad[v].push_back({u, w});
    }
    solve(1);
    ans *= 2;
    cout << ans;
}

Compilation message (stderr)

Main.cpp: In function 'int main()':
Main.cpp:164:16: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
  164 |         freopen(task".inp", "r", stdin);
      |         ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~
Main.cpp:165:16: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
  165 |         freopen(task".out", "w", stdout);
      |         ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...