Submission #1234518

#	Time	Username	Problem	Language	Result	Execution time	Memory
1234518		nai0610	Winter Driving (CCO19_day1problem3)	C++20	25 / 25	97 ms	31560 KiB

Main

#include <bits/stdc++.h>
#define ll int64_t
#define ld long double
using namespace std;
// 
const int maxn =2e5+5;
const int mod = 1e9+7; // 998244353,1610612741
const ll inf = 1e18;
const ld pi = atan(1.0L)*4;
int n;
ll a[maxn],sz[maxn],mx[maxn];
vector<int> g[maxn];
void dfs(int u,int p=-1) {
	for (auto v:g[u]){
		if (v!=p) {
			dfs(v,u);
			mx[u]=max(mx[u],sz[v]);
		}
	}
}
int get(int u,int p=-1) {
	sz[u]=a[u];
	for (auto v:g[u]) {
		if (v!=p) sz[u]+=get(v,u);
	}
	return sz[u];
}
void calc(int u,int p,ll& s,ll d){
	s+=a[u]*d;
	for (auto v:g[u]) {
		if (v!=p) calc(v,u,s,d+a[u]);
	}

}
vector<ll> get(int l,int r,const vector<ll>& v) {
	int m=r-l+1;
	vector<ll> res;
	for (int mask=0;mask<(1<<m);mask++) {
		ll s=0;
		for (int i=0;i<m;i++) {
			if (mask&(1<<i)) s+=v[i+l];
		}
		res.push_back(s);
	}
	sort(res.begin(),res.end());
	return res;
}
int main() {
    // freopen("../input.inp","r",stdin);
    // freopen("output.out","w",stdout);
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    clock_t start,end;
	start=clock();
	cin >>n;
	ll s=0,s2=0;
	for (int i=1;i<=n;i++) cin >>a[i],s+=a[i],s2+=a[i]*(a[i]-1);
	for (int i=2;i<=n;i++) {
		int p;
		cin >>p;
		g[p].push_back(i);
		g[i].push_back(p);
	}
	get(1);
	dfs(1);
	vector<int> cen;
	for (int u=1;u<=n;u++){
		if (max(mx[u],s-sz[u])*2<=s) cen.push_back(u);
	}
	ll res=0;
	for (auto u:cen){
		get(u);
		vector<ll> b;
		ll temp=0;
		for (auto v:g[u]) {
			calc(v,u,temp,a[u]);
			b.push_back(sz[v]);
		}
		int m=b.size();
		int mid=m/2;
		vector<ll> l=get(0,mid-1,b),r=get(mid,m-1,b);
		ll val=0,S=s-a[u];
		for (auto i:l) {
			if (i>S/2) continue;
			int pos=upper_bound(r.begin(),r.end(),S/2-i)-r.begin()-1;
			if (pos>=0&&r[pos]+i<=S/2) val=max(val,r[pos]+i);
		}
		res=max(res,temp+val*(S-val));
	}
	cout <<res+s2;
	end=clock();
	ld dur=(ld)(end-start)/(ld)(CLOCKS_PER_SEC)*(1000.0L);
	cerr<<dur<<'\n';
    return 0;
}

• Subtask 1: 5 / 5
• Subtask 2: 5 / 5
• Subtask 3: 5 / 5
• Subtask 4: 5 / 5
• Subtask 5: 5 / 5