PTA 特立独行的幸福（带权并查集）

L2-029 特立独行的幸福 (25 分)

传送门

输入样例 1

10 40

输出样例 1

19 8
23 6
28 3
31 4
32 3

注意

样例中，10、13 也都是幸福数，但它们分别依附于其他数字（如 23、31 等等），所以不输出。其它数字虽然其实也依附于其它幸福数，但因为那些数字不在给定区间 [10, 40] 内，所以它们在给定区间内是特立独行的幸福数。

输入样例 2

110 120

输出样例 2

SAD

思路

每个数字迭代后只能产生唯一一个数字，并且需要求“依附于它的的幸福数的个数”（距离祖先的距离），可以用带权并查集来做。
fa[maxn]记录父结点，dis[maxn]记录距离父结点的距离，End[maxn]记录是否为“特立独行”。add(a,b)将a的父结点设置为b，若a在给定区间内，则标记b为“不独立”。
并查集不可以有环，加入一条边时要判断。

代码

#include<bits/stdc++.h>
using namespace std;
const int maxn=1e4+5;
int fa[maxn],dis[maxn],End[maxn];
int l,r,has;

void add(int a,int b){
    fa[a]=b;dis[a]=1;
    if(l<=a&&a<=r)  End[b]=1;
}

int get(int x){
    if(fa[x]==0)    return x;
    int root=get(fa[x]);
    dis[x]+=dis[fa[x]];
    return fa[x]=root;
}

void dfs(int x){
    int t=x,ans=0;
    while(t){
        int tmp=t%10;
        ans+=tmp*tmp;
        t/=10;
    }
    if(get(ans)!=get(x))    add(x,ans);
    if(!fa[ans]&&ans!=1)    dfs(ans);
}

bool isprime(int x){
    if(x<2) return false;
    for(int i=2;i*i<=x;i++){
        if(x%i==0)  return false;
    }
    return true;
}

int main(){
    cin>>l>>r;
    for(int i=l;i<r;i++)    dfs(i);
    for(int i=l;i<r;i++){
        if(!End[i]&&get(i)==1){
            has=true;
            cout<<i<<" ";
            if(isprime(i))  printf("%d\n",2*dis[i]);
            else    printf("%d\n",dis[i]);
        }
    }
    if(!has)    puts("SAD");
}