Apples and Snakes#12242

Apples and Snakes (Çözgüdiň düşündirişi)

Problem Author:	Andrii Kholod (Connect)
Prepared by:	Andrii Kholod (Connect)
Editorial:	Andrii Kholod (Connect)

Solution

We remove all vertices that contain snakes. After deleting the forbidden vertices, we obtain a forest. We need to maximize the sum of the weights of unique vertices in the union of no more than two simple paths. Two paths together can form up to four leaves — hence the combinations $t o p$ - $2/3/4$ .

For each “clean” component, we run a $D F S$ and compute for each vertex $v$ :

$d p_{v} =$ dp_v = the best downward branch starting at $v$ (one endpoint is $v$ ),
$m x_p a t h_{v} =$ mx\_path_v = the best single path within the subtree of $v$ ,
$u p_{v} =$ up_v = the best combination of two paths in the subtree of $v$ , where at least one passes through or ends at $v$ .

Combinations through $v$ (taking the largest values among its children):

$d p_{v} = a_{v} + max (d p_{c hi l d})$ dp_v = a_v + \max(dp_{child}),
$m x_p a t h_{v} = max (m x_p a t h_{v}, a_{v} + top-2 d p_{c hi l d})$ mx\_path_v = \max(mx\_path_v, a_v + \text{top-2 } dp_{child}),
$u p_{v} = max (u p_{v}, a_{v} + top-3 d p_{c hi l d})$ up_v = \max(up_v, a_v + \text{top-3 } dp_{child}),
$an s = max (an s, a_{v} + top-4 d p_{c hi l d})$ ans = \max(ans, a_v + \text{top-4 } dp_{child})

We also merge different subtrees using prefix maximums:

$an s = max (an s, max (m x_p a t h_{c hi l d 1} + m x_p a t h_{c hi l d 2}, u p_{c hi l d} + a_{v} + d p_{o t h er}, a_{v} + d p_{x} + u p_{x}))$ ans = \max(ans, \max(mx\_path_{child1} + mx\_path_{child2},\ up_{child} + a_v + dp_{other},\ a_v + dp_x + up_x))

Additionally, we combine a “ready” path from one subtree with branches from others, ensuring we never reuse the same child twice.

Finally, for each component we have $m x_p a t h_{r oo t}$ mx\_path_{root}. We take the two largest of them and update once more:

$an s = max (an s, m x_p a t h_{r oo t_{1}} + m x_p a t h_{r oo t_{2}})$ ans = \max(ans, mx\_path_{root_1} + mx\_path_{root_2})

The asymptotic complexity of this solution is $O (\sum d e g (v) lo g (d e g (v)))$ O\left(\sum deg(v) \log(deg(v))\right).

#include <bits/stdc++.h>

//#pragma GCC optimize("Ofast", "unroll-loops")
//#pragma GCC target("sse", "sse2", "sse3", "ssse3", "sse4")

#include <bits/stdc++.h>

#define all(a) a.begin(),a.end()
#define len(a) (int)(a.size())
#define mp make_pair
#define pb push_back
#define fi first
#define se second

using namespace std;

typedef pair<int, int> pii;
typedef long long ll;
typedef long double ld;
template<class T>
using vec = vector<T>;

template<typename T>
bool umin(T &a, T b) {
    if (b < a) { a = b; return true; }
    return false;
}
template<typename T>
bool umax(T &a, T b) {
    if (a < b) { a = b; return true; }
    return false;
}

#ifdef KoRoVa
#define DEBUG for (bool _FLAG = true; _FLAG; _FLAG = false)
#define LOG(...) print(#__VA_ARGS__" ::", __VA_ARGS__) << endl
template <class ...Ts> auto &print(Ts ...ts) { return ((cerr << ts << " "), ...); }
#else
#define DEBUG while (false)
#define LOG(...)
#endif

const int max_n = 2e5 + 11, inf = 1000111222;
const ll linf = inf * 1ll * inf;

int a[max_n];
int bad[max_n];
vector <int> edge[max_n];
int used[max_n];
ll dp[max_n];
ll mx_path[max_n];
ll up[max_n];
ll ans = 0;

inline void dfs (int v, int p = -1) {
    used[v] = 1;
    ll pref = 0, pref_up = 0, mx_pref = 0;
    vector <pair<ll, int>> one, cur;

    for (int to : edge[v]) {
        if (to == p || bad[to]) continue;
        dfs(to, v);
        umax(dp[v], dp[to]);
        umax(mx_path[v], mx_path[to]);
        cur.pb({dp[to], to});

        umax(up[v], up[to] + a[v]);

        if (umax(ans, up[to] + a[v] + pref)) { LOG(ans); }
        if (umax(ans, a[v] + pref_up + dp[to])) { LOG(ans, v, pref_up, dp[to]); }

        umax(pref, dp[to]);
        umax(pref_up, up[to]);
        umax(ans, mx_path[to] + mx_pref);

        umax(mx_pref, mx_path[to]);
        one.pb({mx_path[to], to});
    }

    dp[v] += a[v];
    sort(all(cur));
    reverse(all(cur));
    ll sum = a[v], now = a[v], mx = a[v];

    for (int i = 0; i < len(cur) && i < 4; i++) {
        if (i < 3) sum += cur[i].first;
        if (i < 2) mx += cur[i].first;
        now += cur[i].first;
    }

    umax(ans, now);
    umax(up[v], sum);
    umax(mx_path[v], mx);

    sort(all(one));
    reverse(all(one));

    for (int i = 0; i < len(one) && i < 3; i++) {
        umax(up[v], a[v] + one[i].first);
        for (int j = 0; j < 3 && j < len(cur); j++) {
            if (cur[j].second == one[i].second) continue;
            umax(up[v], a[v] + one[i].first + cur[j].first);
            for (int k = 0; k < 3 && k < len(cur); k++) {
                if (cur[k].second == one[i].second || j == k) continue;
                if (umax(ans, a[v] + one[i].first + cur[j].first + cur[k].first)) {
                    LOG(v, ans);
                }
            }
        }
    }

    if (umax(ans, up[v])) { LOG(v, ans); }
}

int main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0);

    int n, m;
    cin >> n >> m;
    for (int i = 0; i < n; i++) cin >> a[i];
    for (int i = 0; i < m; i++) {
        int x; cin >> x; --x; bad[x] = 1;
    }
    for (int i = 1, u, v; i < n; i++) {
        cin >> u >> v; --u; --v;
        edge[u].pb(v);
        edge[v].pb(u);
    }

    vector<ll> gg;
    for (int i = 0; i < n; i++) {
        if (bad[i] || used[i]) continue;
        dfs(i);
        gg.pb(mx_path[i]);
    }

    sort(all(gg));
    reverse(all(gg));
    ll cur = 0;
    for (int i = 0; i < len(gg) && i < 2; i++) cur += gg[i];
    umax(ans, cur);

    cout << ans << '\n';
}

/*
KoRoVa!
*/

C++