Submission #240549

#TimeUsernameProblemLanguageResultExecution timeMemory
240549NucleistChase (CEOI17_chase)C++14
100 / 100
1635 ms258536 KiB
#include <bits/stdc++.h> 
using namespace std; 
#pragma GCC target ("avx2")
#pragma GCC optimization ("O3")
#pragma GCC optimization ("unroll-loops")
#define flash ios_base::sync_with_stdio(false); cin.tie(0); cout.tie(0)
#define debug(x) cerr << " - " << #x << ": " << x << endl;
#define debugs(x, y) cerr << " - " << #x << ": " << x << " " << #y << ": " << y << endl;
#define all(x) (x).begin(),(x).end()
#define sz(x) (ll)x.size()
#define ll long long
#define INF 1000000000
#define MOD 1000000007
#define pb push_back
#define ve vector<ll>
#define dos pair<ll,ll>
#define vedos vector<dos>
#define rand mt19937 rng(chrono::steady_clock::now().time_since_epoch().count())
#define EPS 0.000001
struct greateri
{
    template<class T>
    bool operator()(T const &a, T const &b) const { return a > b; }
};
void setIO(string s) {
  ios_base::sync_with_stdio(0); cin.tie(0); 
  freopen((s+".in").c_str(),"r",stdin);
  freopen((s+".out").c_str(),"w",stdout);
}
ll n,v;
ve tree[200101];
set<pair<ll,int>>ka[101];
ll ni[100001],best[100001][101],best1[100001][101][2];
int w[100001];
int pa[100001];
vector<int> eul;
ll kol=0;
int dep[100001];
int kal;
void dfs(int node,int p){
  pa[node]=p;
  for(auto it:tree[node]){
    if(it!=p){
      kal++;
      dfs(it,node);
      kal--;
    }
  }
  dep[node]=kal;
  eul.pb(node);
}
ll solve(ll node,ll p){
  ll ans=0;
  for (int i = 0; i <= v; ++i)
  {
    ka[i].clear();
  }
  for(auto it:tree[node]){
    if(it!=p){
      ka[0].insert({0,it});
    }
  }
  int indx=0;
  for (ll i = 1; i <= v; ++i)
  {
    best1[node][i][0]=0;
    best1[node][i][1]=0;
    for(auto it:tree[node]){
    if(it!=p){
        ll cur=ni[node];
        ll cur1=cur;
        if(p!=-1)cur-=w[p];
        cur1-=w[it];
        best1[node][i][0]=max(best1[node][i][0],cur+best1[it][i-1][0]);
        best1[node][i][0]=max(best1[node][i][0],best1[it][i][0]);
        best1[node][i][1]=max(best1[node][i][1],cur1+best1[it][i-1][1]);
        best1[node][i][1]=max(best1[node][i][1],best1[it][i][1]);
        if(sz(ka[i])>1){
          auto u=ka[i].end();
          u--;
          u--;
          if((*u).first<best1[it][i][0])ka[i].insert({best1[it][i][0],it});
        }
        else ka[i].insert({best1[it][i][0],it});
      }
    }
    indx=1;
    best1[node][i][1]=max(best1[node][i][1],best1[node][i-1][1]);
    best1[node][i][0]=max(best1[node][i][0],best1[node][i-1][0]);
    best1[node][i][1]=max(best1[node][i][1],ni[node]);
  } 
  //if(i)best1[node][i-1][0]=max(best1[node][i-1][0],ni[node]);
  for (ll i = 0; i <= v; ++i)
  {
    for(auto it:tree[node]){
    if(it!=p){
        ll yo=best1[it][i][1];
        best[node][v]=max(best[node][v],yo);
        auto u=ka[v-i].end();
        if(ka[v-i].size()>0){
          u--;
          if((*u).second!=it)yo+=(*u).first;
          else if(sz(ka[v-i])>1){
            u--;yo+=(*u).first;
          }
        }
        best[node][v]=max(best[node][v],yo);
 
        if(i!=v){
          yo=0;
          yo+=ni[node];
          yo+=best1[it][i][1];
          yo-=w[it];
          best[node][v]=max(best[node][v],yo);
          auto u=ka[v-i-1].end();
          if(ka[v-i-1].size()>0){
            u--;
            if((*u).second!=it){
              yo+=(*u).first;
            }
            else if(sz(ka[v-i-1])>1){
              u--;
              yo+=(*u).first;
            }
          }
          best[node][v]=max(best[node][v],yo);
        } 
      }
    }
  }
  //debugs(best1[node][v][1],best1[node][v][0])
  return max(ans,best[node][v]);
}
int main()
{
  flash;
  //freopen("tasi.txt","r",stdin);
  //freopen("out","w",stdout);
  cin>>n>>v;
  for (ll i = 0; i < n; ++i)
  {
    cin>>w[i];
  }
  for (ll i = 0; i < n-1; ++i)
  {
    ll u,v;
    cin>>u>>v;
    u--,v--;
    tree[u].pb(v),tree[v].pb(u);
  }
  for (ll i = 0; i < n; ++i)
  {
    for(auto it:tree[i]){
      ni[i]+=w[it];
    }
  }
  ll ans=0;
  dfs(0,-1);
  vedos hed;
  for (int i = 0; i < n; ++i)
  {
    hed.pb({dep[i],i});
  }
  sort(all(hed));
  reverse(all(hed));
  ll ko=hed[0].second;
  for(auto it:eul){
    ans=max(ans,solve(it,pa[it]));
  }
  eul.clear();
  dfs(ko,-1);
  for (int i = 0; i < n; ++i)
  {
    best[i][v]=0;
  }
  for(auto it:eul){
    ans=max(ans,solve(it,pa[it]));
  }
  cout<<ans;
  return 0;
}
//code the AC sol !
// BS/queue/map

Compilation message (stderr)

chase.cpp:4:0: warning: ignoring #pragma GCC optimization [-Wunknown-pragmas]
 #pragma GCC optimization ("O3")
 
chase.cpp:5:0: warning: ignoring #pragma GCC optimization [-Wunknown-pragmas]
 #pragma GCC optimization ("unroll-loops")
 
chase.cpp: In function 'long long int solve(long long int, long long int)':
chase.cpp:63:7: warning: variable 'indx' set but not used [-Wunused-but-set-variable]
   int indx=0;
       ^~~~
chase.cpp: In function 'void setIO(std::__cxx11::string)':
chase.cpp:27:10: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)', declared with attribute warn_unused_result [-Wunused-result]
   freopen((s+".in").c_str(),"r",stdin);
   ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~
chase.cpp:28:10: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)', declared with attribute warn_unused_result [-Wunused-result]
   freopen((s+".out").c_str(),"w",stdout);
   ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...