Submission #535986

#	Time	Username	Problem	Language	Result	Execution time	Memory
535986		Killer2501	Janjetina (COCI21_janjetina)	C++14	110 / 110	366 ms	19308 KiB

Main

#include <bits/stdc++.h>
#define ll long long
#define ld long double
#define ull unsigned long long
#define pb push_back
#define pll pair<ll, ll>
#define pii pair<int, int>
#define fi first
#define se second
using namespace std;
const int N = 2e5+2;
const int M = 52;
const int mod = 1e9+7;
const ll base = 75;
const ll inf = 1e16;
int n, lab[N], t, c[N], m, k, a[N], b[N], d[N], tong;
ll ans;
vector<pii> adj[N];
vector<int> vi, cur;
bool vis[N];
int r, sz;
ll pw(ll k, ll n)
{
    ll total = 1;
    for(; n; n >>= 1)
    {
        if(n&1)total = total * k % mod;
        k = k * k % mod;
    }
    return total;
}
void add(int& x, int y)
{
    x += y;
    if(x >= mod)x -= mod;
}
int findp(int u)
{
    return lab[u] < 0 ? u : lab[u] = findp(lab[u]);
}
bool hop(int u, int v)
{
    u = findp(u);
    v = findp(v);
    if(u == v)return false;
    if(lab[u] > lab[v])swap(u, v);
    lab[u] += lab[v];
    lab[v] = u;
    return true;
}
struct BIT
{
    int n;
    vector<ll> fe;
    BIT(){}
    void init(int _n)
    {
        n = _n;
        fe.resize(n+1, 0);
    }
    void reset()
    {
        fill(fe.begin(), fe.end(), 0);
    }
    void add(int id, int x)
    {
        if(id <= 0)exit(0);
        for(; id <= n; id += id & -id)fe[id] += x;
    }
    ll get(int id)
    {
        ll total = 0;
        id = min(id, n);
        if(id <= 0)return total;
        for(; id; id -= id & -id)total += fe[id];
        return total;
    }
}bit;
void dfs(int u, int p = 0)
{
    d[u] = 1;
    int mx = 0;
    for(pii x: adj[u])
    {
        int v = x.se;
        if(v == p || b[v])continue;
        dfs(v, u);
        d[u] += d[v];
        mx = max(mx, d[v]);
    }
    mx = max(mx, sz-d[u]);
    if(mx*2 <= sz)r = u;
}
bool cmp(const int& x, const int& y)
{
    return (a[x] < a[y]) || (a[x] == a[y] && c[x] < c[y]);
}
void cal(int u, int p = 0)
{
    for(pii v: adj[u])
    {
        if(v.se == p || b[v.se])continue;
        a[v.se] = max(a[u], v.fi);
        c[v.se] = c[u]+1;
        cal(v.se, u);
    }
    cur.pb(u);
}
void fix()
{
    vector<int> res;
    int i = 0, j = 0;
    while(i < (int) cur.size() || j < (int) vi.size())
    {
        if(i == (int) cur.size())
        {
            res.pb(vi[j]);
            ++j;
        }
        else if(j == (int) vi.size() || cmp(cur[i], vi[j]))
        {
            res.pb(cur[i]);
            ++i;
        }
        else
        {
            res.pb(vi[j]);
            ++j;
        }
    }
    swap(res, vi);
    cur.clear();
}
void decom(int u)
{
    int all = sz;
    b[u] = 1;
    //cout << u <<'\n';
    for(pii x: adj[u])
    {
        if(b[x.se])continue;
        a[x.se] = x.fi;
        c[x.se] = 1;
        cal(x.se);
        sort(cur.begin(), cur.end(), cmp);
        for(int v: cur)
        {
            ans -= bit.get(a[v]-c[v]);
            bit.add(c[v], 1);
        }
        for(int v: cur)bit.add(c[v], -1);
        fix();
    }
    for(int v: vi)
    {
        //if(u == 1)cout << a[v] <<" "<<c[v]<<" "<<v<<'\n';
        ans += bit.get(a[v]-c[v]);
        bit.add(c[v], 1);
    }
    for(int v: vi)
    {
        bit.add(c[v], -1);
        if(a[v] >= c[v])++ans;
    }
    //cout << u <<" "<<ans<<'\n';
    vi.clear();
    for(pii x: adj[u])
    {
        int v = x.se;
        if(b[v])continue;
        sz = (d[v] > d[u] ? all - d[u] : d[v]);
        dfs(v);
        decom(r);
    }
}
void sol()
{
    cin >> n >> k;
    for(int i = 1; i < n; i ++)
    {
        int x, y, w;
        cin >> x >> y >> w;
        w -= k;
        adj[x].pb({w, y});
        adj[y].pb({w, x});
    }
    bit.init(n);
    sz = n;
    dfs(1);
    decom(r);
    cout << ans*2;
}
int main()
{
    cin.tie(0);
    cout.tie(0);
    ios_base::sync_with_stdio(0);
    #define task "test"
    if(fopen(task".inp", "r"))
    {
        freopen(task".inp", "r", stdin);
        freopen(task".out", "w", stdout);
    }
    int test = 1;
    //cin >> test;
    while(test -- > 0)sol();
    return 0;
}
/*
1234
21
*/

Compilation message (stderr)

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

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