HDU 4812 D Tree (点分治，技巧去同一子树的贡献)

D Tree

Problem Description

There is a skyscraping tree standing on the playground of Nanjing University of Science and Technology. On each branch of the tree is an integer (The tree can be treated as a connected graph with N vertices, while each branch can be treated as a vertex). Today the students under the tree are considering a problem: Can we find such a chain on the tree so that the multiplication of all integers on the chain (mod 106 + 3) equals to K?
Can you help them in solving this problem?

Input

There are several test cases, please process till EOF.
Each test case starts with a line containing two integers N(1 <= N <= 105) and K(0 <=K < 106 + 3). The following line contains n numbers vi(1 <= vi < 106 + 3), where vi indicates the integer on vertex i. Then follows N - 1 lines. Each line contains two integers x and y, representing an undirected edge between vertex x and vertex y.

Output

For each test case, print a single line containing two integers a and b (where a < b), representing the two endpoints of the chain. If multiply solutions exist, please print the lexicographically smallest one. In case no solution exists, print “No solution”(without quotes) instead.
For more information, please refer to the Sample Output below.

Sample Input

5 60
2 5 2 3 3
1 2
1 3
2 4
2 5
5 2
2 5 2 3 3
1 2
1 3
2 4
2 5

Sample Output

3 4
No solution

Hint

1. “please print the lexicographically smallest one.”是指: 先按照第一个数字的大小进行比较，若第一个数字大小相同，则按照第二个数字大小进行比较，依次类推。

2. 若出现栈溢出，推荐使用C++语言提交，并通过以下方式扩栈：
   #pragma comment(linker,"/STACK:102400000,102400000")

题意： 一棵树，每一个节点有一个值 val[i]，输出两个端点 a, b 。表示 a 到 b 的路径上的点权积 % mod == k，存在多个则输出字典序最小。 其中 mod = 1e6 + 3

思路：

此处全是废话：有没有大佬这样写A了的，教教我。。 像一般的点分治一样，再 cal 的时候传参确定这个 cal 是用来加还是用来减的。 对于所有的 dis 都放入 vector(包括距离，以及 id). 然后排序先按距离由小到大，其次是 id 由大到小。 每一次都 upper_bound() 另一半距离，然后 - 1 的位置如果符合 乘积 % mod == k，那么就把这两个点 hash 一下，放入 unordered_map，计算贡献(是加还是减)，最后遍历整个 unordered_map

正解：

普通点分治，处理 dis 的时候有点技巧。每一次从 重心处 往儿子处理 dis，每当处理完一棵子树后再更新 flag (flag[i] 表示为到重心距离 i 的最小节点编号)，这样就不会出现两个距离出现在同一棵子树上的情况。
dis 的时候就直接判断当前的距离是否 == k，或者是 flag[另一半] 是否出现过，直接更新答案即可。
每当处理完一个重心以后，再清空 flag 。

Code：

#pragma comment(linker,"/STACK:102400000,102400000")
#include "bits/stdc++.h"
#define debug(x) cerr << "[" << #x <<": " << (x) <<"]"<< endl
#define pii pair<int,int>
#define clr(a, b) memset((a),b,sizeof(a))
#define rep(i, a, b) for(int i = a;i < b;i ++)
#define pb push_back
#define MP make_pair
#define LL long long
#define ull unsigned LL
#define ls i << 1
#define rs (i << 1) + 1
#define fi first
#define se second
#define ptch putchar
#define CLR(a) while(!(a).empty()) a.pop()
using namespace std;

inline LL read() {
    LL s = 0, w = 1;
    char ch = getchar();
    while (!isdigit(ch)) {
        if (ch == '-')
            w = -1;
        ch = getchar();
    }
    while (isdigit(ch))
        s = s * 10 + ch - '0', ch = getchar();
    return s * w;
}

const int maxn = 1e5 + 10;
const int inf = 0x3f3f3f3f;
const int mod = 1e6 + 3;
LL inv[mod],k;
int head[maxn], cnt;
int sonnum[maxn],val[maxn],n;
bool vis[maxn];
int flag[mod],e[maxn],e2[maxn],id[maxn],p1,p2;
int a,b;

struct xx {
    int v, nex;
} edge[maxn << 1];

LL pow(LL a, LL b, LL p) {  //快速幂 a^n % p
    LL ans = 1;
    while(b) {
        if(b & 1)
            ans = ans * a % p;
        a = a * a % p;
        b >>= 1;
    }
    return ans;
}

LL niYuan(LL a, LL b) { //费马小定理求逆元
    return pow(a, b - 2, b);
}

void init() {
    for (int i = 0; i <= n + 5; ++i) {
        head[i] = -1;
        vis[i] = 0;
    }
    cnt = 0;
}

void Tree_center(int &mi, int &rt, int p, int fa, int sum) {  /// 处理重心
    int maxx = 0;
    sonnum[p] = 1;
    for (int i = head[p]; ~i; i = edge[i].nex) {
        int v = edge[i].v;
        if (v == fa || vis[v]) continue;
        Tree_center(mi, rt, v, p, sum);
        sonnum[p] += sonnum[v];
        maxx = max(maxx,sonnum[v]);
    }
    if (mi > max(maxx, sum - sonnum[p])) {
        mi = max(maxx, sum - sonnum[p]);
        rt = p;
    }
}

void updateAns(int x,int y){    /// 更新答案
    if(x == y) return ;
    if(x > y) swap(x,y);
    if(a > x || x == a && b > y){
        a = x;
        b = y;
    }
}

void Tree_dis(int p, int fa, LL d,int rt) {
    if(d == k){              /// 当前一个链的距离就是 k 
        updateAns(p,rt);
    }
    LL fx = k * val[rt] * inv[d] % mod;
    if(flag[fx]){          /// 另一半是否出现
        updateAns(flag[fx],p);
    }
    e[++ p1] = d;
    e2[++ p2] = d;
    id[p1] = p;
    for (int i = head[p]; ~i; i = edge[i].nex) {
        int v = edge[i].v;
        if (vis[v] || v == fa) continue;
        Tree_dis(v, p, (d * 1LL * val[v]) % mod,rt);
    }
}

void dfs(int p,int fa) {
    int mi = inf,rt;
    Tree_center(mi,rt,p,fa,sonnum[p]);
    vis[rt] = 1;
    p2 = 0;
    for(int i = head[rt]; ~i;i = edge[i].nex){
        int v = edge[i].v;
        if(vis[v]) continue;				///这句话能够优化 1s
        p1 = 0;
        Tree_dis(v,rt,1LL * val[rt] * val[v] % mod,rt);
        for(int j = 1;j <= p1;++ j){     /// dis 完一棵子树之后再更新 flag
            if(!flag[e[j]]) flag[e[j]] = id[j];
            else flag[e[j]] = min(flag[e[j]],id[j]);
        }
    }
    for(int i = 1;i <= p2;++ i) flag[e2[i]] = 0;   /// 重新清空 flag，然后递归向下一个重心
    for(int i = head[rt]; ~i;i = edge[i].nex){
        int v = edge[i].v;
        if(vis[v] || v == fa) continue;
        dfs(v,rt);
    }
}

int main() {
    inv[0] = inv[1] = 1;
    for (int i = 2; i < mod; i++) inv[i] = niYuan(i,mod);
    while (~scanf("%d%lld", &n, &k)) {
        init();
        for (int i = 1; i <= n; ++i) {
            val[i] = read();
        }
        for (int i = 1; i < n; ++i) {
            int u = read(), v = read();
            edge[cnt] = xx{v, head[u]};
            head[u] = cnt++;
            edge[cnt] = xx{u, head[v]};
            head[v] = cnt++;
        }
        a = inf, b = 0;
        sonnum[1] = n;
        dfs(1, -1);
        if (a != inf) printf("%d %d\n", a, b);
        else printf("No solution\n");
    }
    return 0;
}