HDU 4670 Cube number on a tree

divide and conquer on tree.

 

#include <map>

#include <vector>

#include <cstdio>

#include <cstdlib>

#include <cstring>

#include <iostream>

using namespace std;

const int N = 5e4+10;

const int K = 32;

typedef long long LL;



struct edge_t {

  int v, to;

};

vector<edge_t> E;

int L[N];

void init(int n) {

  E.clear();

  memset(L, -1, sizeof(int)*n);

}

void add_edge(int u, int v) {

  edge_t t = {v, L[u]};

  L[u] = E.size();

  E.push_back(t);

}



int n, k;

LL w[N], p[K];



LL ll_split(LL x) {

  LL rv = 0;

  for (int i = 0; i < k; i++) {

    LL o = 0;

    for ( ; x%p[i] == 0; x/= p[i]) o++;

    rv |= o%3<<(i<<1);

  }

  return rv;

}

LL ll_add(LL lhs, LL rhs) {

  LL rv = 0;

  for (int i = 0; i < k; i++) {

    LL o = ((lhs>>(i<<1)&3)+(rhs>>(i<<1)&3))%3;

    rv |= o<<(i<<1);

  }

  return rv;

}

LL ll_com(LL x) {

  LL rv = 0;

  for (int i = 0; i < k; i++) {

    rv |= (3-(x>>(i<<1)&3))%3<<(i<<1);

  }

  return rv;

}

void ll_show(LL x) {

  for (int i = 0; i < k; i++) {

    printf("%lld ", x>>(i<<1)&3);

  }

}



bool vis[N];

int size, s[N];

void get_size(int u) {

  vis[u] = 1;

  s[u] = 1;

  for (int i = L[u]; ~i; i = E[i].to) {

    int v = E[i].v;

    if (vis[v]) continue;

    get_size(v);

    s[u] += s[v];

  }

  vis[u] = 0;

}

int core;

void find_core(int u) {

  vis[u] = 1;

  int flag = 1;

  for (int i = L[u]; ~i; i = E[i].to) {

    int v = E[i].v;

    if (vis[v]) continue;

    if (s[v] > size/2) flag = 0;

    find_core(v);

  }

  if (flag && size-s[u] <= size/2) core = u;

  vis[u] = 0;

}

vector<LL> dis;

void get_dis(int u, LL d) {

  vis[u] = 1;

  dis.push_back(ll_add(w[u], d));

  for (int i = L[u]; ~i; i = E[i].to) {

    int v = E[i].v;

    if (vis[v]) continue;

    get_dis(v, ll_add(w[u], d));

  }

  vis[u] = 0;

}

int dac(int u) {

  get_size(u);

  size = s[u];

  if (size == 1) return !w[u];

  find_core(u);

  int rv = !w[u = core];

  vis[u] = 1;

  map<LL, int> m;

  m[0]++;

  for (int i = L[u]; ~i; i = E[i].to) {

    int v = E[i].v;

    if (vis[v]) continue;

    dis.clear();

    get_dis(v, 0);

    for (int j = 0; j < dis.size(); j++) {

      map<LL, int>::iterator it = m.find(ll_com(ll_add(dis[j], w[u])));

      if (it != m.end()) rv += it->second;

    }

    for (int j = 0; j < dis.size(); j++) m[dis[j]]++;

  }

  for (int i = L[u]; ~i; i = E[i].to) {

    int v = E[i].v;

    if (vis[v]) continue;

    rv += dac(v);

  }

  vis[u] = 0;

  return rv;

}



int main() {

  for ( ; ~scanf("%d", &n); ) {

    init(n);

    scanf("%d", &k);

    for (int i = 0; i < k; i++) {

      char b[32];

      scanf("%s", b);

      p[i] = atoll(b);

    }

    for (int i = 0; i < n; i++) {

      char b[32];

      scanf("%s", b);

      w[i] = ll_split(atoll(b));

    }

    for (int i = 0; i < n-1; i++) {

      int u, v;

      scanf("%d%d", &u, &v);

      u--, v--;

      add_edge(u, v);

      add_edge(v, u);

    }

    printf("%d\n", dac(0));

  }

  return 0;

}