Автор задачі: | Андрій Холод (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!
*/