Submission #859444

#	Time	Username	Problem	Language	Result	Execution time	Memory
859444		lucri	Paths (RMI21_paths)	C++17	100 / 100	239 ms	20536 KiB

Main

#include <bits/stdc++.h>
using namespace std;
long long dp[200010],sum,n,x,y,z,k,ans[100010],or1[100010],or2[100010];
set<pair<long long,long long>>s;
vector<vector<pair<int,int>>>a;
set<pair<long long,long long>>::iterator it,it2;
void initializeaza(int nod,int t)
{
    for(auto x:a[nod])
    {
        if(x.first==t)
        {
            dp[nod]=x.second;
            continue;
        }
        initializeaza(x.first,nod);
    }
    int pozmax=0;
    for(auto x:a[nod])
    {
        if(x.first==t)
            continue;
        if(dp[pozmax]<=dp[x.first])
            pozmax=x.first;
    }
    dp[nod]+=dp[pozmax];
    for(auto x:a[nod])
    {
        if(x.first==t||x.first==pozmax)
            continue;
        s.insert({dp[x.first],x.first});
    }
}
void swaprad(int rad,int newrad,int val)
{
    if(or1[rad]==0)
    {
        for(auto x:a[rad])
        {
            if(dp[x.first]>=dp[or1[rad]])
            {
                or1[rad]=x.first;
            }
        }
    }
    if(or2[rad]==0)
    {
        for(auto x:a[rad])
        {
            if(x.first!=or1[rad]&&dp[x.first]>=dp[or2[rad]])
            {
                or2[rad]=x.first;
            }
        }
    }
    int leg=or1[rad];
    if(leg==newrad)
        leg=or2[rad];
    it2=s.find({dp[leg],leg});
    if(*it2>=*it)
    {
        --it;
        sum-=dp[leg];
        sum+=(*it).first;
    }
    s.erase({dp[leg],leg});
    dp[rad]=dp[leg]+val;
    s.insert({dp[rad],rad});
    it2=s.find({dp[rad],rad});
    if(*it2>*it)
    {
        sum-=(*it).first;
        ++it;
        sum+=dp[rad];
    }
    it2=s.find({dp[newrad],newrad});
    if(*it2>=*it)
    {
        --it;
        sum-=dp[newrad];
        sum+=(*it).first;
    }
    s.erase({dp[newrad],newrad});
    dp[newrad]=0;
    if(or1[newrad]==0)
    {
        for(auto x:a[newrad])
        {
            if(dp[x.first]>=dp[or1[newrad]])
            {
                or1[newrad]=x.first;
            }
        }
    }
    if(or2[newrad]==0)
    {
        for(auto x:a[newrad])
        {
            if(x.first!=or1[newrad]&&dp[x.first]>=dp[or2[newrad]])
            {
                or2[newrad]=x.first;
            }
        }
    }
    leg=or1[newrad];
    if(leg==rad)
        leg=or2[newrad];
    s.insert({dp[leg],leg});
    it2=s.find({dp[leg],leg});
    if(*it2>*it)
    {
        sum-=(*it).first;
        ++it;
        sum+=dp[leg];
    }
    ans[newrad]=sum;
}
void raspunde(int nod,int nodant,int val)
{
    for(auto x:a[nod])
    {
        if(x.first==nodant||x.first>n)
            continue;
        swaprad(nod,x.first,x.second);
        raspunde(x.first,nod,x.second);
    }
    if(nodant)
        swaprad(nod,nodant,val);
}
int main()
{
    ios_base::sync_with_stdio(false);
    cin.tie(0);
    cout.tie(0);
    cin>>n>>k;
    a.resize(2*n+5);
    for(int i=1;i<n;++i)
    {
        cin>>x>>y>>z;
        a[x].push_back({y,z});
        a[y].push_back({x,z});
    }
    int nr=n;
    for(int i=1;i<=n;++i)
    {
        if(a[i].size()==1)
        {
            ++nr;
            a[i].push_back({nr,0});
            a[nr].push_back({i,0});
        }
    }
    initializeaza(1,0);
    for(auto x:a[1])
    {
        if(s.find({dp[x.first],x.first})==s.end())
        {
            or1[1]=x.first;
            s.insert({dp[x.first],x.first});
            break;
        }
    }
    dp[1]=0;
    if(k>=s.size())
    {
        for(auto x:s)
            sum+=x.first;
        for(int i=1;i<=n;++i)
            cout<<sum<<'\n';
        return 0;
    }
    it=s.end();
    for(int i=1;i<=k;++i)
    {
        --it;
        sum+=(*it).first;
    }
    ans[1]=sum;
    raspunde(1,0,0);
    for(int i=1;i<=n;++i)
        cout<<ans[i]<<'\n';
    return 0;
}

Compilation message (stderr)

Main.cpp: In function 'int main()':
Main.cpp:164:9: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::set<std::pair<long long int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  164 |     if(k>=s.size())
      |        ~^~~~~~~~~~

• Subtask 1: 8 / 8
• Subtask 2: 11 / 11
• Subtask 3: 17 / 17
• Subtask 4: 20 / 20
• Subtask 5: 12 / 12
• Subtask 6: 32 / 32