Submission #1244828

#	Time	Username	Problem	Language	Result	Execution time	Memory
1244828		CrabCNH	Janjetina (COCI21_janjetina)	C++20	110 / 110	271 ms	22204 KiB

Main

#include <bits/stdc++.h>

#define task     "BriantheCrab"

#define int    long long
#define pii    pair <int, int>
#define fi     first
#define se     second
#define szf    sizeof
#define sz(s)  (int)((s).size())
#define all(v) (v).begin(), (v).end()

typedef long long ll;
typedef unsigned long long ull;
typedef long double ld;

using namespace std;

template <class T> void minimize (T &t, T f) {if (t > f) t = f;}
template <class T> void maximize (T &t, T f) {if (t < f) t = f;}

const int maxN = 2e5 + 5;
const int inf = 1e18 + 7;
const int mod = 1e9 + 7;

int n, k;
int sz[maxN];
bool del[maxN];
vector<pii> adj[maxN];
vector<pii> paths;
int res;

struct fenwick {
    int bit[maxN * 2];
    fenwick () {
        memset (bit, 0, sizeof (bit));
    }
    
    void update (int x, int val) {
        for (; x < maxN; x += x & -x) {
            bit[x] += val;
        }
    }
    
    int get (int x) {
        int res = 0;
        for (; x > 0; x -= x & -x) {
            res += bit[x];
        }
        return res;
    }
} T;

int dfsSz (int u, int p) {
    sz[u] = 1;
    for (auto [v, w] : adj[u]) {
        if (v == p || del[v]) {
            continue;
        }
        sz[u] += dfsSz (v, u);
    }
    return sz[u];
}

int findCen (int u, int p, int num) {
    for (auto [v, w] : adj[u]) {
        if (v == p || del[v]) {
            continue;
        }
        if (sz[v] > num / 2) {
            return findCen (v, u, num);
        }
    }
    return u;
}

void dfs (int u, int p, int dep, int val) {
    paths.push_back ({val, dep});
    for (auto [v, w] : adj[u]) {
        if (v == p || del[v]) {
            continue;
        }
        dfs (v, u, dep + 1, max (val, w));
    }
}

int cal (vector <pii> &Path) {
    sort (all (Path));
    int curAns = 0;
    for (auto [w, len] : Path) {
        if (w >= len) {
            curAns += T.get (w - len);
        }
        T.update (k + len, 1);
    }
    for (auto [w, len] : Path) {
        T.update (k + len, -1);
    }
    return curAns;
}

void demp (int u) {
    u = findCen (u, 0, dfsSz (u, 0));
    del[u] = true;
    paths.clear ();
    dfs (u, 0, 0, 0);
    res += cal (paths);
    for (auto [v, w] : adj[u]) {
        if (!del[v]) {
            paths.clear ();
            dfs (v, u, 1, w);
            res -= cal (paths);
        }
    }
    for (auto [v, w] : adj[u]) {
        if (!del[v]) {
            demp (v);
        }
    }
}

void solve () {
    cin >> n >> k;
    for (int i = 0; i < n - 1; i ++) {
        int u, v, w;
        cin >> u >> v >> w;
        adj[u].push_back ({v, w});
        adj[v].push_back ({u, w});
    }
    demp (1);
    cout << res * 2 << '\n';
    return;
}

signed main () {
    cin.tie (nullptr) -> sync_with_stdio (false);
    if (fopen (task".inp", "r")) {
        freopen (task".inp", "r", stdin);
        freopen (task".out", "w", stdout);
    }
    int t = 1;
    while (t --) {
        solve ();
    }
    return 0;
}
// thfv

Compilation message (stderr)

Main.cpp: In function 'int main()':
Main.cpp:138:17: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
  138 |         freopen (task".inp", "r", stdin);
      |         ~~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~
Main.cpp:139:17: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
  139 |         freopen (task".out", "w", stdout);
      |         ~~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~

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