Cod sursa(job #2466938)

Utilizator AlexPop28Pop Alex-Nicolae AlexPop28 Data 3 octombrie 2019 11:59:02
Problema A+B Scor 0
Compilator cpp-64 Status done
Runda teme_upb Marime 2.01 kb
#include <bits/stdc++.h>
#define DEBUG(x) cerr << (#x) << ": " << x << endl

using namespace std;

const int MAXBIT = 31;
const int MOD = (int)1e9 + 7;

struct DSU {
  struct Node {
    int dad, sz;
    vector<int> info;
    Node() : dad(-1), sz(1), info(MAXBIT) {}
  };
  DSU(int n_, const vector<int> &v) : n(n_), dsu(n) {
    for (int i = 0; i < n; ++i) {
      for (int bit = 0; bit < MAXBIT; ++bit) {
        dsu[i].info[bit] += (v[i] & (1 << bit)) != 0;
      }
    }
  }
  int Find(int node) {
    return dsu[node].dad == -1 ? node : (dsu[node].dad = Find(dsu[node].dad));
  }
  void Union(int x, int y) {
    x = Find(x), y = Find(y);
    if (x == y) return;
    if (dsu[x].sz < dsu[y].sz) swap(x, y);
    dsu[y].dad = x;
    dsu[x].sz += dsu[y].sz;
    for (int bit = 0; bit < MAXBIT; ++bit) {
      dsu[x].info[bit] += dsu[y].info[bit];
    }
  }
  int n;
  vector<Node> dsu;
};

void DFS(int node, int parent, vector<vector<int>> &adj) {
  for (int &x : adj[node]) {
    adj[x].erase(find(adj[x].begin(), adj[x].end(), node));
    DFS(x, node, adj);
  }
}


int main() {
  //ifstream cin("countfefete.in");
  //ofstream cout("countfefete.out");

  int n; cin >> n;
  vector<int> v(n);
  for (int &x : v) cin >> x;
  vector<vector<int>> adj(n);
  for (int i = 0; i < n - 1; ++i) {
    int x, y; cin >> x >> y; --x, --y;
    adj[x].emplace_back(y);
    adj[y].emplace_back(x);
  }
  DFS(0, -1, adj);
  vector<int> order(n);
  iota(order.begin(), order.end(), 0);
  sort(order.begin(), order.end(), [&](int a, int b) {
    return val[a] > val[b];  
  });
  vector<int> pow2(n + 1);
  pow2[0] = 1;
  for (int i = 1; i <= n; ++i) {
    pow2[i] = pow2[i - 1] + pow2[i - 1];
    if (pow2[i] >= MOD) pow2[i] -= MOD;
  }
  DSU dsu(n, val);
  vector<bool> vis(n);
  for (int it = 0; it < n; ++it) {
    int node = order[it];
    vis[node] = true;
    for (int &x : adj[node]) {
      if (!vis[x]) continue;
      dsu.Union(node, x);
    }
     
  }
}