HDU 4670 Cube number on a tree

Description

The country Tom living in is famous for traveling. Every year, many tourists from all over the world have interests in traveling there. 

There are n provinces in the country. According to the experiences from the tourists came before, every province has its own preference value. A route’s preference value from one province to another is defined as the product of all the preference value of the provinces on the route. It’s guaranteed that for each two provinces in the country there is a unique route from one to another without passing any province twice or more. 

Tom is a boy crazy about cube number. A cube number is a positive integer whose cube root is also an integer. He is planning to travel from a province to another in the summer vacation and he will only choose the route with the cube number preference value. Now he want to know the number of routes that satisfy his strange requirement.

Input

The input contains several test cases, terminated by EOF. 

Each case begins with a number n ( 1 ≤ n ≤ 50000), the number of the provinces. 

The second line begins with a number K (1 ≤ K ≤ 30), and K difference prime numbers follow. It’s guaranteed that all the preference number can be represented by the product of some of this K numbers(a number can appear multiple times). 

The third line consists of n integer numbers, the ith number indicating the preference value P 

i(0 ≤ P 

i ≤ 10 

15) of the i-th province. 

Then n - 1 lines follow. Each line consists of two integers x, y, indicating there is a road connecting province x and province y. 

Output

For each test case, print a number indicating the number of routes that satisfy the requirement.

Sample Input

5

3 2 3 5

2500 200 9 270000 27

4 2

3 5

2 5

4 1

Sample Output

1

#include<map>
#include<cmath>
#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
typedef long long LL;
const int INF = 0x7FFFFFFF;
const int mod = 1e9 + 7;
const int maxn = 1e5 + 10;
int n, m, a[30];
LL x, y;

struct point
{
  int c[30];
  point() { memset(c, 0, sizeof(c)); }
  bool operator<(const point &a) const
  {
    for (int i = 0; i < m; i++)
    {
      if (c[i] < a.c[i]) return true;
      if (c[i] > a.c[i]) return false;
    }
    return false;
  }
  bool operator==(const point &a) const
  {
    for (int i = 0; i < m; i++) if (c[i] != a.c[i]) return false;
    return true;
  }
};

point rev(point b)
{
  point a;
  for (int i = 0; i < m; i++) a.c[i] = (3 - b.c[i]) % 3;
  return a;
}
map<point, int> M;
map<point, int>::iterator it;

point operator +(const point &a, const point &b)
{
  point c;
  for (int i = 0; i < m; i++) c.c[i] = (a.c[i] + b.c[i]) % 3;
  return c;
}

struct Tree
{
  int ft[maxn], nt[maxn], u[maxn], sz;
  int vis[maxn], mx[maxn], ct[maxn];
  point v[maxn];
  void clear(int n)
  {
    mx[sz = 0] = INF;
    for (int i = 1; i <= n; i++) ft[i] = -1, vis[i] = 0;
  }
  void AddEdge(int x, int y)
  {
    u[sz] = y;  nt[sz] = ft[x]; ft[x] = sz++;
    u[sz] = x;  nt[sz] = ft[y]; ft[y] = sz++;
  }
  int dfs(int x, int fa, int sum)
  {
    int y = mx[x] = (ct[x] = 1) - 1;
    for (int i = ft[x]; i != -1; i = nt[i])
    {
      if (vis[u[i]] || u[i] == fa) continue;
      int z = dfs(u[i], x, sum);
      ct[x] += ct[u[i]];
      mx[x] = max(mx[x], ct[u[i]]);
      y = mx[y] < mx[z] ? y : z;
    }
    mx[x] = max(mx[x], sum - ct[x]);
    return mx[x] < mx[y] ? x : y;
  }
  LL get(int x, int fa, point y)
  {
    LL ans = M[rev(y)];
    for (int i = ft[x]; i != -1; i = nt[i])
    {
      if (u[i] == fa || vis[u[i]]) continue;
      ans += get(u[i], x, y + v[u[i]]);
    }
    return ans;
  }
  void put(int x, int fa, point y)
  {
    ++M[y];
    for (int i = ft[x]; i != -1; i = nt[i])
    {
      if (u[i] == fa || vis[u[i]]) continue;
      put(u[i], x, y + v[u[i]]);
    }
  }
  LL find(int x)
  {
    M.clear();  M[point()] = 1;
    LL ans = v[x] == rev(v[x]);
    for (int i = ft[x]; i != -1; i = nt[i])
    {
      if (vis[u[i]]) continue;
      ans += get(u[i], x, v[x] + v[u[i]]);
      put(u[i], x, v[u[i]]);
    }
    return ans;
  }
  LL work(int x, int sum)
  {
    int y = dfs(x, -1, sum);
    LL ans = find(y); vis[y] = 1;
    for (int i = ft[y]; i != -1; i = nt[i])
    {
      if (vis[u[i]]) continue;
      ans += work(u[i], ct[u[i]] > ct[y] ? sum - ct[y] : ct[u[i]]);
    }
    return ans;
  }
}solve;

int main()
{
  while (scanf("%d%d", &n, &m) != EOF)
  {
    solve.clear(n);
    for (int i = 0; i < m; i++) scanf("%d", &a[i]);
    for (int i = 1; i <= n; i++)
    {
      scanf("%lld", &x);
      for (int j = 0, k; j < m; j++)
      {
        for (k = 0; x%a[j] == 0; k++) x /= a[j];
        solve.v[i].c[j] = k % 3;
      }
    }
    for (int i = 1; i < n; i++)
    {
      scanf("%lld%d", &x, &y);
      solve.AddEdge(x, y);
    }
    printf("%d\n", solve.work(1, n));
  }
  return 0;
}