Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- #include <bits/stdc++.h>
- #define F first
- #define S second
- #define pb push_back
- using namespace std;
- typedef long long ll;
- typedef pair <int, int> pii;
- #include <ext/pb_ds/assoc_container.hpp>
- #include <ext/pb_ds/tree_policy.hpp>
- using namespace __gnu_pbds;
- #define ordered_set tree<pair<ll,int>, null_type,less<pair<ll,int>>, rb_tree_tag,tree_order_statistics_node_update>
- const int INF = 1e8;
- const int mxN = 1e5 + 9;
- const int mxM = 5e4 + 9;
- const int MOD = 998244353;
- vector <int> dfsOrder, adj[mxN];
- int dp[mxN][1 << 7], r[mxN], sz[mxN][7], d[mxN];
- void goDFS (int v, int p) {
- dfsOrder.pb (v);
- sz[v][d[v]] = 1;
- for (auto u: adj[v]) {
- if (u == p)
- continue;
- goDFS (u, v);
- for (int i = 0; i <= 5; i++)
- sz[v][i] += sz[u][i];
- }
- r[v] = dfsOrder.size() - 1;
- return;
- }
- void solve () {
- int n;
- cin >> n;
- for (int i = 1; i <= n; i++)
- cin >> d[i];
- int two;
- cin >> two;
- for (int i = 1; i < n; i++) {
- int u, v;
- cin >> u >> v;
- adj[u].pb (v);
- adj[v].pb (u);
- }
- goDFS (1, -1);
- dfsOrder.push_back (0);
- for (int i = 0; i < mxN; i++)
- for (int j = 0; j < (1 << 7); j++)
- dp[i][j] = -INF;
- dp[dfsOrder.size() - 1][0] = 0;
- for (int i = dfsOrder.size() - 2; i >= 0; i--) {
- int v = dfsOrder [i];
- for (int msk = 0; msk < (1 << 7); msk++) {
- dp[i][msk] = dp[i + 1][msk];
- int vMask = 0, total = 0;
- for (int j = 0; j <= 5; j++) {
- vMask |= ((sz[v][j] & 1) << j);
- total += sz[v][j];
- }
- int prevMask = msk ^ vMask;
- dp[i][msk] = max (dp[i][msk], dp[r[v] + 1][prevMask] + total);
- }
- }
- int res = dp[0][0];
- for (int i = 0; i <= 5; i++)
- res = max (res, dp[0][1 << i]);
- cout << res << "\n";
- return;
- }
- int main () {
- int t = 1;
- //cin >> t;
- //preCalc ();
- while (t--) {
- //initialize common variables
- //go solve
- solve ();
- }
- return 0;
- }
- /*
- *
- *
- *
- *
- 12
- 5 1 3 2 2 1 4 4 4 3 4 3
- 2
- 1 2
- 1 3
- 2 4
- 2 5
- 3 6
- 3 7
- 6 8
- 6 9
- 7 10
- 7 11
- 10 12
- *
- */
Advertisement
Add Comment
Please, Sign In to add comment