cf1627 Round #766 Div2-C【图论+DFS】

Date：2022.01.17
题意：给定一个树，给每个边加权能不能构造出任意<=2个边相加都是质数。能输出对应边的方案，不能输出-1。

思路：首先单独一个边是质数很好办，每个边都是质数即可。在此前提下，怎么才能保证每两个边相加都是质数？不难发现，一旦存在度数>=3的点就不能满足，因为质数只有偶质数2和奇质数，任意两个奇数相加为偶数且一定是合数，只有一个奇质数和一个偶质数相加才能保证和为质数，因此不存在任意一个点连着三条边且这三条边两两加和都为质数的情况。由此，每个点的度数只能<=2，又是树，满足情况的只能是单链。建个图，dfs每次给边权依次赋值2、3，每次顺序记录边，输出再离线做一下。
代码如下：

#include <iostream>
#include <cstring>
#include <algorithm>
#include <map>
#include <set>
using namespace std;
const int N = 2e5+10;//注意点最多为边的两倍，多开一倍
typedef long long LL;
typedef pair<int, int> PII;
map<PII,LL>q;set<LL>s;
LL a[N],b[N];
LL d[N],h[N],idx,e[N],ne[N];
bool st[N];
LL n,m,t,k;
void add(LL a,LL b)
{
    e[idx]=b;
    ne[idx]=h[a];
    h[a]=idx++;
}
void dfs(LL a,LL ans,LL idx)
{
    if(idx==s.size()) return;
    st[a]=true;
    //cout<<a<<' ';
    for(int i=h[a];i!=-1;i=ne[i])
    {
        if(st[e[i]]) continue;
        PII t1={a,e[i]},t2={e[i],a};
        q[t1]=ans;q[t2]=ans;
        if(ans==2) dfs(e[i],3,idx+1);
        else if(ans==3) dfs(e[i],2,idx+1);
    }
}
int main()
{
    ios::sync_with_stdio(false);cin.tie(0);cout.tie(0);
    cin>>t;
    while(t--)
    {
        bool flag=true;
        memset(d,0,sizeof d);memset(h,-1,sizeof h);memset(st,0,sizeof st);
        q.clear();s.clear();
        cin>>n;
        for(int i=1;i<=n-1;i++)
        {
            cin>>m>>k;a[i]=m;b[i]=k;
            d[m]++;d[k]++;s.insert(m);s.insert(k);
            add(m,k);add(k,m);
        }
        for(auto it:s) 
            if(d[it]>2) {flag=false;break;}
        if(!flag) {cout<<"-1"<<endl;continue;}
        else
        {
            for(auto it:s)
                if(d[it]==1) {dfs(it,2,0);break;}
            for(int i=1;i<=n-1;i++)
            {
                PII t={a[i],b[i]};
                cout<<q[t]<<' ';
            }
            cout<<endl;
        }
    }
    return 0;
}