【算法练习】HDU - 4812 D Tree （点分治）

版权声明：本文为博主原创文章，未经博主允许不得转载。 https://blog.csdn.net/pengwill97/article/details/82563054

题意

求树上是否有路径点权成绩mod1e6+7 = k。若有多对路径，输出节点编号最小的一组。

题解

逆元 + 点分治即可。

代码

#include<bits/stdc++.h>
using namespace std;
typedef double db;
typedef long long ll;
typedef unsigned long long ull;
const int nmax = 1e6+7;
const int INF = 0x3f3f3f3f;
const ll LINF = 0x3f3f3f3f3f3f3f3f;
const ull p = 67;
const ll mod = 1e6 + 3;
int head[100005], tot;
struct edge {
    int to, nxt;
}e[100005<<1];
void add_edge(int u, int v) {
    e[tot].to = v;
    e[tot].nxt = head[u];
    head[u] = tot++;
}
ll val[100005];
int n, k, sz[100005], root, totsubtree, nowmx;
bool visit[100005];
int inv[nmax], ss[100005], id[100005], ts, top;
int ansx, ansy;
int val_of_map[nmax], times_of_map[nmax];
void getinv() {
    inv[0] = inv[1] = 1;
    for(int i = 2; i <= mod; ++i) {
        inv[i] = 1ll * (mod - mod / i) * inv[mod % i] % mod;
    }
}
void insert_map(int key, int val) {
    if(times_of_map[key] == ts) {
        if(val_of_map[key] > val)
            val_of_map[key] = val;
    } else {
        times_of_map[key] = ts;
        val_of_map[key] = val;
    }

}
void getroot(int u, int f) {
    sz[u] = 1;
    int mxchild = 0;
    for(int i = head[u]; i != -1; i = e[i].nxt) {
        int v = e[i].to;
        if(v != f && !visit[v]) {
            getroot(v, u);
            sz[u] += sz[v];
            mxchild = max(mxchild, sz[v]);
        }
    }
    int tmp = max(mxchild, totsubtree - sz[u]);
    if(tmp < nowmx) {
        nowmx = tmp;
        root = u;
    }
}
void dfs(int u, int f, ll d) {
    ss[top] = d;
    id[top++] = u;
    for(int i = head[u]; i != -1; i = e[i].nxt) {
        int v = e[i].to;
        if(v != f && !visit[v]) {
            dfs(v, u, (d * val[v]) % mod);
        }
    }
}
void update_ans(int a, int b) {
    if(a > b) swap(a, b);
    if(a < ansx) {
        ansx = a;
        ansy = b;
    } else if(a == ansx && b < ansy) {
        ansy = b;
    }
}
void solve(int u) {
    visit[u] = true;
    ts ++;
    for(int i = head[u]; i != -1; i = e[i].nxt) {
        int v = e[i].to;
        if(visit[v]) continue;
        top = 0;
        dfs(v, u, val[v] % mod);
        for(int j = 0; j < top; ++j) {
            int p = (val[u] * ss[j]) % mod, q = (1ll * k * inv[p]) % mod;
            if(p == k) update_ans(u, id[j]);
            if(times_of_map[q] == ts) update_ans(val_of_map[q], id[j]);

        }
        for(int j = 0; j < top; ++j)
            insert_map(ss[j], id[j]);
    }

    for(int i = head[u]; i != -1; i = e[i].nxt) {
        int v = e[i].to;
        if(visit[v]) continue;
        totsubtree = sz[v]; nowmx = INF; root = 0;
        getroot(v, u);
        solve(root);
    }
}
int main() {
    getinv();
    while(scanf("%d %d", &n, &k)!= EOF) {
        memset(head, -1, sizeof head);
        ts = tot = 0;
        ansx = ansy = nowmx = INF;
        memset(visit, 0, sizeof visit);
        memset(sz, 0, sizeof sz);
        memset(val_of_map, 0, sizeof val_of_map);
        memset(times_of_map, 0, sizeof times_of_map);
        for(int i = 1; i <= n; ++i)
            scanf("%lld", &val[i]);
        for(int i = 1; i <= n - 1; ++i) {
            int u, v;
            scanf("%d %d", &u, &v);
            add_edge(u, v);
            add_edge(v, u);
        }
        totsubtree = n;
        getroot(1, -1);
        solve(root);
        if(ansx == INF) {
            printf("No solution\n");
        } else {
            printf("%d %d\n", ansx, ansy);
        }
    }
    return 0;
}