Half Nice Years Gym - 101840H （点分治 or 并查集）

题目

　　https://cn.vjudge.net/problem/Gym-101840H

题意

　　给出一棵树，问有多少对点，将他们之间的边权相乘之后所获得的值仅有两个不同的质因子。

题解

　　当时没想着用并查集做，写了个点分治= =。我们对于重心的子树挨个搜索，用number数组记录下质因子数量为0， 1， 2的点的数量，对于质因子数量为1的查询时还要知道多少与你当前想要链接的点为相同的质因子，所以用has【i】记录下质因子为 i 的点的数量。那么能与质因子数量为1的点相连的点就是number【1】-has【arr【j】】（arr【j】是当前选择的点所拥有的质因子）。

#include <iostream>
#include <cstring>
#include <string>
#include <algorithm>
#include <cmath>
#include <cstdio>
#include <queue>
#include <stack>
#include <map>
#include <bitset>
#define ull unsigned long long
#define met(a, b) memset(a, b, sizeof(a))
#define lowbit(x) (x&(-x))
#define MID (l + r) / 2
#define ll long long

using namespace std;

const int maxn = 1e5 + 7;
const ll mod = 1e6 + 3;
const int inf = 0x3f3f3f3f;
const ll INF = 0x3f3f3f3f3f3f3f3f;

struct Edge {
    int to, dist, net;
}edge[maxn * 2];

int n, k;
int head[maxn], cnt;
int val[maxn];
int dis[maxn], arr[maxn], tail, all, sz[maxn], msz[maxn], rt;
bool vis[maxn];
int res;
bitset<maxn> is_prime;
int prime[maxn], ans;
int number[3];
int has[maxn];
int N[40];    

void num() {
    is_prime[1] = 1;
    for(int i = 2; i < maxn; i++) {
        if(!is_prime[i]) prime[ans++] = i;
        for(int j = 0; j < ans && i * prime[j] <= n; j++) {
            is_prime[i * prime[j]] = 1;
            if(i % prime[j] == 0) break;
        }
    }
}
void init() {
    res = 0;
    for(int i = 0; i <= n; i++) {
        head[i] = -1;
        vis[i] = 0;
    }
    met(has, 0);
    cnt = 0;
}
void addedge(int u, int v, int c) {
    edge[cnt] = (Edge){v, c, head[u]};
    head[u] = cnt++;
}
void Findrt(int pos, int pre) {
    sz[pos] = 1;
    msz[pos] = 0;
    for(int i = head[pos]; i != -1; i = edge[i].net) {
        int to = edge[i].to;
        if(to == pre || vis[to]) continue;
        Findrt(to, pos);
        sz[pos] += sz[to];
        msz[pos] = max(msz[pos], sz[to]);
    }
    msz[pos] = max(msz[pos], all - sz[pos]);
    if(msz[pos] < msz[rt] || rt == 0) rt = pos;
}
int calc(int v) {
    int sum = 0;
    for(int i = 0; i < ans && sum < 3 && prime[i] <= v; i++) {
        int flag = 1;
        while(sum < 3 && v % prime[i] == 0) {
            if(flag) {
                flag = 0;
                N[++sum] = prime[i], v /= prime[i];
            }
            else sum = 3;
        }
    }
    return sum;
}
void dfs(int pos, int pre, int v) {
    dis[pos] = dis[pre] + calc(v);
    if(dis[pos] == 2) arr[++tail] = -1;
    else if(dis[pos] == 1) arr[++tail] = N[1];
    else if(dis[pos] == 0) arr[++tail] = 0;
    for(int i = head[pos]; i != -1; i = edge[i].net) {
        int to = edge[i].to;
        if(vis[to] || to == pre) continue;
        dfs(to, pos, edge[i].dist);
    }
}
void divide(int pos) {
    vis[pos] = 1;
    for(int i = head[pos]; i != -1; i = edge[i].net) {
        int to = edge[i].to;
        if(vis[to]) continue;
        tail = 0;
        dfs(to, pos, edge[i].dist);
        for(int j = 1; j <= tail; j++) {
            if(arr[j] == -1) res += number[0] + 1;
            else if(arr[j] == 0) res += number[2];
            else res += number[1] - has[arr[j]];
        }
        for(int j = 1; j <= tail; j++) {
            if(arr[j] == -1) number[2]++;
            else if(arr[j] == 0) number[0]++;
            else number[1]++, has[arr[j]]++;
        }
    }
　　//原本这里是要递归进子树清除对进入下一层的影响，memset时间复杂度太高。但当时调不出来qwq
    tail = 0;
    met(has, 0);
    met(dis, 0);
    for(int i = 0; i < 3; i++) number[i] = 0;
    for(int i = head[pos]; i != -1; i = edge[i].net) {
        int to = edge[i].to;
        if(vis[to]) continue;
        all = sz[to]; 
        rt = 0;
        Findrt(to, 0);
        divide(rt);
    }
}

int main() {
    freopen("evaluations.in", "r", stdin);
    num();
    int T, k = 0;
    cin >> T;
    while(T--) {
        cin >> n;
        init();
        for(int i = 1; i < n; i++) {
            int u, v, c;
            cin >> u >> v >> c;
            addedge(u, v, c);
            addedge(v, u, c);
        }
        all = n;
        rt = 0;
        Findrt(1, 0);
        divide(rt);
        printf("Case %d: %d\n", ++k, res);
    }
    return 0;
}