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
;
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
, 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
;
int size, s
;
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;
}