Submission #925050

#	Time	Username	Problem	Language	Result	Execution time	Memory
925050		boris_mihov	Power Plant (JOI20_power)	C++17	100 / 100	153 ms	28612 KiB

power

#include <algorithm>
#include <iostream>
#include <numeric>
#include <cassert>
#include <cstring>
#include <vector>
#include <queue>
#include <stack>
#include <set>

typedef long long llong;
const int MAXN = 2e5 + 10;
const int INF  = 2e9;

int n;
int a[MAXN];
int b[MAXN];
int sz[MAXN];
char s[MAXN];
int dp[MAXN][2];
bool bl[MAXN][2];
std::vector <int> g[MAXN];

void buildDFS(int node, int par)
{
    for (int &u : g[node])
    {
        if (u == par)
        {
            std::swap(u, g[node].back());
            g[node].pop_back();
            break;
        }
    }

    sz[node] = s[node] - '0';
    for (const int &u : g[node])
    {
        buildDFS(u, node);
        sz[node] += sz[u];
    }
}

int f(int node, bool parentHas)
{
    if (bl[node][parentHas])
    {
        return dp[node][parentHas];
    }

    bl[node][parentHas] = true;
    dp[node][parentHas] = std::max(dp[node][parentHas], (int)(s[node] == '1'));

    if (!parentHas && s[node] == '1')
    {
        for (const int &u : g[node])
        {
            dp[node][parentHas] = std::max(dp[node][parentHas], f(u, true) + 1);
        }
    }

    if (!parentHas)
    {
        for (const int &u : g[node])
        {
            dp[node][parentHas] = std::max(dp[node][parentHas], f(u, false));
        }
    }

    int curr = 0;
    for (const int &u : g[node])
    {
        curr += f(u, true);    
    }

    dp[node][parentHas] = std::max(dp[node][parentHas], curr - (s[node] == '1'));
    return dp[node][parentHas];
}

void solve()
{
    buildDFS(1, 0);
    std::cout << f(1, 0) << '\n';
}

void input()
{
    std::cin >> n;
    for (int i = 1 ; i < n ; ++i)
    {
        int u, v;
        std::cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }

    std::cin >> s + 1;
}

void fastIOI()
{
    std::ios_base :: sync_with_stdio(0);
    std::cout.tie(nullptr);
    std::cin.tie(nullptr);
}   

int main()
{
    fastIOI();
    input();
    solve();

    return 0;
}

Compilation message (stderr)

power.cpp: In function 'void input()':
power.cpp:97:19: warning: suggest parentheses around '+' inside '>>' [-Wparentheses]
   97 |     std::cin >> s + 1;
      |                 ~~^~~

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