wwbt3

Posted 2022-06-09 znsbc-13

#include<bits/stdc++.h>
using namespace std;
#define ll int
#define A 3000000
ll n,K,pf[A],a[A],last[A],pos[A],colo[A];
bool arkn[A];
vector<ll> son[350000];
bool color(ll x,ll r,ll l,ll col,ll pre)
    colo[x]=col;
    for(ll i=0;i<son[x].size();i++)
        ll y=son[x][i];
        if(y>r||y<l||y==pre) 
            continue;
        
        if(!colo[y])
            colo[y]=colo[x]^1;
            if(!color(y,r,l,3-col,x))
                return 0;
        
        if(colo[y]==colo[x])
            return 0;
    
    return 1;

ll main()

    scanf("%d%d",&n,&K);
    for(ll i=1;i<=512;i++)
         pf[i]=i*i,arkn[i*i]=1;
    for(ll i=1;i<=n;i++) 
        scanf("%d",&a[i]);
    if(K==1)
        ll ans=1,p=n;
        for(ll i=n;i>=1;i--)
            for(ll j=512;j>=1;j--)
                if(pf[j]<=a[i]) break;
                if(last[pf[j]-a[i]])
                    for(ll k=p;k>i;k--) 
                        last[a[k]]=0;
                    pos[ans]=i,p=i,ans++;
                    break;
                
            
            last[a[i]]++;
        
        printf("%d\n",ans);
        for(ll i=ans-1;i>0;i--) 
            printf("%d ",pos[i]);
        puts("");
        return 0;
    
    else
        ll ans=1,p=n;
        for(ll i=1;i<=512;i++)
            arkn[i*i]=1;
        
        for(ll i=n;i>=1;i--)
            for(ll j=p;j>i;j--)
                if(arkn[a[i]+a[j]])
                    son[i].push_back(j);
                    son[j].push_back(i);
                
            
            if(!color(i,p,i,1,0))
                pos[ans]=i,p=i,ans++;
                son[i].clear();
            
            for(ll j=p;j>=i;j--)
                colo[j]=0;
            
        
        printf("%d\n",ans);
        for(ll i=ans-1;i>0;i--)
            printf("%d ",pos[i]);
        puts("");
        return 0;
    

View Code