Яблука та змії#12242

Яблука та змії (Розбір)

Автор задачі:

Андрій Холод (Connect)

Підготовлено:

Андрій Холод (Connect)

Розбір:

Андрій Холод (Connect)

Рішення

Видалимо вершини зі зміями. Після видалення заборонених вершин маємо ліс. Потрібно максимізувати суму ваг унікальних вершин в об’єднанні не більше двох простих шляхів. Два шляхи разом дають до чотирьох листків — звідси комбінації top-2/3/4.

Для кожної «чистої» компоненти запускаємо і рахуємо для вершини :

  • dp_v найкраща гілка вниз від (один кінець у ),

  • _path_v найкращий один шлях у піддереві ,

  • up_v найкраще об’єднання двох шляхів у піддереві , де хоча б один проходить/закінчується у .

Комбінації через (беремо найбільші значення серед дітей):

  • dp_v a_v + (dp_{child}),

  • _path_v (_path_v , a_v + top-2 dp_{child}),

  • up_v (up_v , a_v + top-3 dp_{child}),

  • = ( , a_v , + top-4 dp_{child})

Додатково зшиваємо різні піддерева префіксними максимумами:

  • = ( , (_path_{child1} + _path_{child2} , up_{child} + a_v + dp_{other} , a_v + dp_{x} + up_x));

Також комбінуємо «готовий» шлях з одного піддерева з гілками інших, не використовуючи одну й ту ж дитину двічі.

Наприкінці для кожної компоненти маємо _path_{root}. Беремо дві найбільші та ще раз оновлюємо:

  • ( , _path_{root_1} + _path_{root_2})

Асимптотика такого рішення (deg(v)(deg(v))).

#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() {
//    freopen("object.txt", "r", stdin);
//    freopen("output.txt", "w", stdout);

    ios_base::sync_with_stdio(0);
    cin.tie(0);

    /// all cases
    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]);
        LOG(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';
//    for (int i = 0; i < (1 << n); i++) {
//        ll v = 0;
//        for (int j = 0; j < n; j++) {
//            if (i & (1 << j)) {
//                v += a[j];
//            }
//        }
//        if (v == ans) {
//            LOG(i);
//            for (int j = 0; j < n; j++) {
//                if (i & (1 << j)) {
//                    cout << j + 1 << '\n';
//                }
//            }
//        }
//    }

}

/*
KoRoVa!
*/