Submission #751185

#	Time	Username	Problem	Language	Result	Execution time	Memory
751185		bgnbvnbv	Janjetina (COCI21_janjetina)	C++14	110 / 110	809 ms	74100 KiB

Main

#include<bits/stdc++.h>
#define TASKNAME "codeforce"
#define pb push_back
#define pli pair<int,int>
#define fi first
#define se second
#define fastio ios_base::sync_with_stdio(false); cin.tie(NULL);
using namespace std;
using ll=int;
const ll maxN=1e5+10;
const ll inf=1e8;
const ll mod=1e9+7;
ll sz[maxN];
vector<ll>G[maxN];
bool r[maxN];
int dfs(int u, int p) {
  sz[u]=1;
  for(int v:G[u])
  {
      if(r[v]==false&&v!=p)
      {
          sz[u]+=dfs(v,u);
      }
  }
  return sz[u];
}
int centroid(int u, int p, int siz) {
    for(int v:G[u])
    {
        if(v!=p&&r[v]==false)
        {
            if(sz[v]>siz/2) return centroid(v,u,siz);
        }
    }
    return u;
}
ll pa[maxN];

void build(int u, int p) {
    ll C=centroid(u,u,dfs(u,u));
    r[C]=true;
    pa[C]=p;
    //adj[p].pb(C);
    for(int v:G[C])
    {
        if(r[v]==false)
        {
            build(v,C);
        }
    }
}
vector<ll>val;
ll bit[maxN];
void update(ll i)
{
    while(i<val.size()+5)
    {
        bit[i]++;
        i+=i&(-i);
    }
}
ll get(ll i)
{
    ll c=0;
    while(i>0)
    {
        c+=bit[i];
        i-=(i&(-i));
    }
    return c;
}
ll n,k;
vector<pli>g[maxN];
ll P[maxN][18],mx[maxN][18],h[maxN];
void DFS(ll u=1,ll p=1)
{
    P[u][0]=p;
    for(int i=1;i<=17;i++) P[u][i]=P[P[u][i-1]][i-1];
    for(int i=1;i<=17;i++)
    {
        mx[u][i]=max(mx[u][i-1],mx[P[u][i-1]][i-1]);
    }
    for(auto vv:g[u])
    {
        int v=vv.fi;
        int w=vv.se;
        if(v!=p)
        {
            h[v]=h[u]+1;
            mx[v][0]=w;
            DFS(v,u);
        }
    }
}
struct qq
{
    ll dis,mx;
    bool operator<(const qq&o)
    {
        return mx<o.mx;
    }
};
ll lca(ll u,ll v)
{
    if(h[u]<h[v]) swap(u,v);
    for(int i=17;i>=0;i--) if(h[u]-(1<<i)>=h[v]) u=P[u][i];
    if(u==v) return u;
    for(int i=17;i>=0;i--)
    {
        if(P[u][i]!=P[v][i]) u=P[u][i],v=P[v][i];
    }
    return P[u][0];
}
ll dis(ll u,ll v)
{
    return h[u]+h[v]-2*h[lca(u,v)];
}
ll maxx(ll u,ll v)
{
    ll k=h[u]-h[v];
    ll aa=-inf;
    for(int i=17;i>=0;i--)
    {
        if(k>>i&1) aa=max(aa,mx[u][i]),u=P[u][i];
    }
    return aa;
}
ll get(ll u,ll v)
{
    ll x=lca(u,v);
    return max(maxx(u,x),maxx(v,x));
}
vector<qq>vec[maxN],vec2[maxN];
void solve()
{
    cin >> n >> k;
    for(int i=1;i<n;i++)
    {
        ll u,v,w;
        cin >> u >> v >> w;
        g[u].pb({v,w});
        g[v].pb({u,w});
        G[u].pb(v);
        G[v].pb(u);
    }
    DFS();
    build(1,0);
    for(int i=1;i<=n;i++)
    {
        ll u=i;
        while(pa[u]!=0)
        {
            qq a={dis(i,pa[u]),get(i,pa[u])};
            vec[pa[u]].pb(a);
            vec2[u].pb(a);
            u=pa[u];
        }
    }
    for(int i=1;i<=n;i++) vec[i].pb({0,-inf}),sort(vec[i].begin(),vec[i].end()),sort(vec2[i].begin(),vec2[i].end());
    long long ans=0;
    for(int i=1;i<=n;i++)
    {
        val.clear();
        val.pb(-1);
        for(int j=0;j<vec[i].size();j++)
        {
            val.pb(vec[i][j].dis);
        }
        sort(val.begin(),val.end());
        for(int j=0;j<vec[i].size();j++)
        {
            ll pos=upper_bound(val.begin(),val.end(),vec[i][j].mx-vec[i][j].dis-k)-val.begin()-1;
            ans+=get(pos);
            pos=lower_bound(val.begin(),val.end(),vec[i][j].dis)-val.begin();
            update(pos);
        }
        fill(bit,bit+val.size()+5,0);
    }
    for(int i=1;i<=n;i++)
    {
        swap(vec[i],vec2[i]);
        val.clear();
        val.pb(-1);
        for(int j=0;j<vec[i].size();j++)
        {
            val.pb(vec[i][j].dis);
        }
        sort(val.begin(),val.end());
        for(int j=0;j<vec[i].size();j++)
        {
            ll pos=upper_bound(val.begin(),val.end(),vec[i][j].mx-vec[i][j].dis-k)-val.begin()-1;
            ans-=get(pos);
            pos=lower_bound(val.begin(),val.end(),vec[i][j].dis)-val.begin();
            update(pos);
        }
        fill(bit,bit+val.size()+5,0);
    }
    cout << ans*2;
}
int main()
{
    fastio
    //freopen(TASKNAME".INP","r",stdin);
    //freopen(TASKNAME".OUT","w",stdout);
    solve();
}

Compilation message (stderr)

Main.cpp: In function 'void update(ll)':
Main.cpp:56:12: warning: comparison of integer expressions of different signedness: 'll' {aka 'int'} and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   56 |     while(i<val.size()+5)
      |           ~^~~~~~~~~~~~~
Main.cpp: In function 'void solve()':
Main.cpp:165:22: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<qq>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  165 |         for(int j=0;j<vec[i].size();j++)
      |                     ~^~~~~~~~~~~~~~
Main.cpp:170:22: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<qq>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  170 |         for(int j=0;j<vec[i].size();j++)
      |                     ~^~~~~~~~~~~~~~
Main.cpp:184:22: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<qq>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  184 |         for(int j=0;j<vec[i].size();j++)
      |                     ~^~~~~~~~~~~~~~
Main.cpp:189:22: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<qq>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  189 |         for(int j=0;j<vec[i].size();j++)
      |                     ~^~~~~~~~~~~~~~

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