HDOJ 4670: Cube number on a tree

题目链接：

http://acm.hdu.edu.cn/showproblem.php?pid=4670

题目大意：

树的每个节点有一个点权，所有的点权都可以被给定的30个质数表示出来。

在树上找合法点对。

合法点对指的是，两点间路径上的所有点（含端点）的点权乘积是立方数的点对。

注意：点对中的两个点可以是相同的，这个坑了我好久，切~

算法：

树的点分治。

每次处理的时候，用一个map保存之前遍历过的子树中的节点到根的路径值，另一个保存当前正在遍历的这棵子树里的路径值。

然后很容易就可以求出以这个根为LCA的合法点对。

PS：这么个题还写了4K，有点儿伤哦。。下回果断搞个FOREACH宏

树分治技能初步get，耶！

代码：

#pragma comment(linker,"/STACK:102400000,102400000")
#include<cstdio>
#include<iostream>
#include<algorithm>
#include<sstream>
#include<cstdlib>
#include<cstring>
#include<string>
#include<climits>
#include<cmath>
#include<queue>
#include<vector>
#include<stack>
#include<set>
#include<map>
#define INF 0x3f3f3f3f
#define eps 1e-8
#define mp make_pair
#define pb push_back
#define st first
#define nd second
using namespace std;

const int MAXN = 50010;
const int MAXM = 30;
typedef long long LL;

int siz[MAXN], maxb[MAXN];
bool vis[MAXN];
vector <int> node;
vector <int> mm[MAXN];
LL tnum[MAXN];
LL prim[MAXM];
LL pow3[MAXM];
map <LL, LL> map1, map2;
typedef map <LL, LL> :: iterator mapit;
int n,k;
LL ans;

inline LL hash(LL x)
{
LL tmp = 0LL;
for (int i = 0; i < k; i ++)
{
int cot = 0;
while (x % prim[i] == 0)
{
x /= prim[i];
cot ++;
}
tmp += (cot % 3) * pow3[i];
}
return tmp;
}

inline LL fadd(LL x, LL y)
{
LL tmp = 0LL;
for (int i = 0 ; i < k; i ++)
{
LL ret1 = (x / pow3[i]) % 3;
LL ret2 = (y / pow3[i]) % 3;
tmp += (ret1 + ret2) % 3 * pow3[i];
}
return tmp;
}

inline LL frev(LL x)
{
LL tmp = 0LL;
for (int i = 0 ; i < k; i ++)
{
LL ret = (x / pow3[i]) % 3;
tmp += ((3 - ret) %3) * pow3[i];
}
return tmp;
}

void pre_dfs(int u, int p)
{
node.pb(u);
maxb[u] = 0;
siz[u] = 1;
for (int i = 0; i < mm[u].size(); i ++)
{
int v = mm[u][i];
if (v == p || vis[v])
{
continue;
}
pre_dfs(v, u);
siz[u] += siz[v];
maxb[u] = max(maxb[u], siz[v]);
}
}

void dfs(int u, int p, LL tmp)
{
mapit it;
(it = map2.find(tmp)) == map2.end() ?
map2[tmp] = 1 : (it -> nd) ++;
for (int i = 0; i < mm[u].size(); i ++)
{
int v = mm[u][i];
if (v == p || vis[v])
{
continue;
}
dfs(v, u, fadd(tmp, tnum[v]));
}
}

void cal(int root)
{
map1[0] = 1LL;
if(! tnum[root])
{
ans ++;
}
for (int i = 0; i < mm[root].size(); i ++)
{
int v = mm[root][i];
if (vis[v])
{
continue;
}
dfs(v, root, tnum[v]);
for (mapit it1 = map2.begin(); it1 != map2.end(); it1 ++)
{
mapit it2 = map1.find(frev(fadd(it1 -> st, tnum[root])));
if (it2 != map1.end())
{
ans += (it1 -> nd) * (it2 -> nd);
}
}
for (mapit it1 = map2.begin(); it1 != map2.end(); it1 ++)
{
mapit it2;
(it2 = map1.find(it1 -> st)) == map1.end() ?
map1[it1 -> st] = it1 -> nd : (it2 -> nd) += (it1 -> nd);
}
map2.clear();
}
map1.clear();
}

void solve(int u)
{
node.clear();
pre_dfs(u, -1);
int num = node.size();
int root, tmp = INT_MAX;
for (int i = 0; i < num; i ++)
{
maxb[node[i]] = max(maxb[node[i]], num - maxb[node[i]] - 1);
if (tmp > maxb[node[i]])
{
tmp = maxb[node[i]];
root = node[i];
}
}
vis[root] = true;
cal(root);
for (int i = 0; i < mm[root].size(); i ++)
{
int v = mm[root][i];
if (!vis[v])
{
solve(v);
}
}
}

int main()
{
pow3[0] = 1LL;
for(int i = 1; i < MAXM; i ++)
{
pow3[i] = pow3[i - 1] * 3;
}
while (scanf("%d %d", &n, &k) == 2)
{
ans = 0LL;
memset(vis, 0, sizeof(vis));
for (int i = 0; i < n; i ++)
{
mm[i].clear();
}
for (int i = 0; i < k; i ++)
{
scanf("%I64d", &prim[i]);
}
for (int i = 0; i < n; i ++)
{
scanf("%I64d", &tnum[i]);
tnum[i] = hash(tnum[i]);
}
for (int i = 1; i < n; i ++)
{
int u, v;
scanf("%d %d", &u, &v);
u --;
v --;
mm[u].push_back(v);
mm[v].push_back(u);
}
solve(0);
printf("%I64d\n", ans);
}
return 0;
}