This submission is migrated from previous version of oj.uz, which used different machine for grading. This submission may have different result if resubmitted.
#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 time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |