Submission #518755

#	Time	Username	Problem	Language	Result	Execution time	Memory
518755		cig32	Power Plant (JOI20_power)	C++17	100 / 100	305 ms	79132 KiB

power

#pragma GCC optimize("Ofast")
#include <bits/stdc++.h>
using namespace std;
mt19937_64 rng((int)std::chrono::steady_clock::now().time_since_epoch().count());
const int MAXN = 3e5 + 10;
const int MOD = 1e9 + 7;
#define int long long
 
vector<int> adj[MAXN];
char c[MAXN];
int dp[MAXN];
int s[MAXN], m[MAXN];
int dpf(int node, int prev) {
  int sum = 0, mx = 0;
  for(int x: adj[node]) {
    if(x != prev) {
      int d = dpf(x, node);
      sum += d;
      mx = max(mx, d);
    }
  }
  s[node] = sum;
  m[node] = mx;
  if(prev == -1 && c[node] == '1') {
    return dp[node] = max(mx + 1, sum - 1);
  }
  if(prev == -1 && c[node] == '0') {
    return dp[node] = sum;
  }
  if(c[node] == '1') {
    return dp[node] = max(1ll, sum - 1);
  }
  else {
    return dp[node] = sum;
  }
}
int ans = 0;

void dfs(int node, int prev) {
  ans = max(ans, dp[node]); 
  vector<int> v;
  for(int x: adj[node]) {
    if(x != prev) v.push_back(dp[x]);
  }
  int w = v.size();
  vector<int> pmax, smax;
  pmax.resize(w);
  smax.resize(w);
  if(w) pmax[0] = v[0], smax[w-1] = v[w-1];
  for(int i=1; i<w; i++) pmax[i] = max(pmax[i-1], v[i]);
  for(int i=w-2; i>=0; i--) smax[i] = max(smax[i+1], v[i]);
  int id = 0;
  for(int x: adj[node]) {
    if(x != prev) {
      int sto1, sto2, sto3, sto4, sto5, sto6;
      sto1 = s[node], sto2 = m[node], sto3 = dp[node];
      sto4 = s[x], sto5 = m[x], sto6 = dp[x];
      //reroot node -> x
      s[node] -= dp[x];
      m[node] = max((id == 0 ? 0 : pmax[id-1]), (id == w-1 ? 0 : smax[id+1]));
      dp[node] = (c[node] == '1' ? max(1ll, s[node] - 1) : s[node]);
      s[x] += dp[node];
      m[x] = max(m[x], dp[node]);
      dp[x] = (c[x] == '1' ? max(m[x] + 1, s[x] - 1) : s[x]);
      dfs(x, node);
      //reroot x -> node
      s[node] = sto1, m[node] = sto2, dp[node] = sto3;
      s[x] = sto4, m[x] = sto5, dp[x] = sto6;
      id++;
    }
  }
}
void solve(int tc) {
  int N;
  cin >> N;
  for(int i=1; i<N; i++) {
    int a, b;
    cin >> a >> b;
    adj[a].push_back(b);
    adj[b].push_back(a);
  }
  if(N == 1) {
    cout << "1\n"; return;
  }
  for(int i=1; i<=N; i++) cin >> c[i];
  dpf(1, -1);
  dfs(1, -1);
  cout << ans << "\n";
}
 
int32_t main(){
  ios::sync_with_stdio(0); cin.tie(0);
  int t = 1; //cin >> t;
  for(int i=1; i<=t; i++) solve(i);
} 

• Subtask 1: 6 / 6
• Subtask 2: 41 / 41
• Subtask 3: 53 / 53