【SPOJ1825】Free tour II （点分治，启发式）

题意：

边权可能为负

思路：

感觉我自己写的还是太过僵硬了，可以灵活一点，比如可以多写几个不同的dfs求出不同的信息，而不是压到同一个dfs里

#include<cstdio>
#include<cstring>
#include<string>
#include<cmath>
#include<iostream>
#include<algorithm>
#include<map>
#include<set>
#include<queue>
#include<vector>
using namespace std;
typedef long long ll;
typedef unsigned int uint;
typedef unsigned long long ull;
typedef pair<int,int> PII;
typedef vector<int> VI;
#define fi first
#define se second s
#define MP make_pair
#define N   410000
#define M   1100000
#define MOD 1000000007
#define eps 1e-8 
#define pi acos(-1)
#define oo 2e9+1

struct node
{
    int x,id;
}b[N];
 

int head[N],vet[N],nxt[N],len[N],dis[N],dep[N],mxdep[N],son[N],flag[N],c[N],a[N],f[N],tmp[N],mx[N],
    n,K,tot,top,ans,sum,root; 
    



int read()
{ 
   int v=0,f=1;
   char c=getchar();
   while(c<48||57<c) {if(c=='-') f=-1; c=getchar();}
   while(48<=c&&c<=57) v=(v<<3)+v+v+c-48,c=getchar();
   return v*f;
}

bool cmp(node a,node b)
{
    return a.x<b.x||
           (a.x==b.x&&a.id<b.id);

}



void add(int a,int b,int c)
{
    nxt[++tot]=head[a];
    vet[tot]=b;
    len[tot]=c;
    head[a]=tot;
}

void getroot(int u,int fa)
{
    son[u]=1; c[u]=0;
    int e=head[u];
    while(e)
    {
        int v=vet[e];
        if(v!=fa&&!flag[v])
        {
            getroot(v,u);
            son[u]+=son[v];
            c[u]=max(c[u],son[v]);
        }
        e=nxt[e];
    }
    c[u]=max(c[u],sum-c[u]);
    if(c[u]<c[root]) root=u;
}

void dfs(int u,int fa)
{
    mxdep[u]=dep[u];
    int e=head[u];
    while(e)
    {
        int v=vet[e];
        if(v!=fa&&!flag[v])
        {
            dep[v]=dep[u]+a[v];
            dis[v]=dis[u]+len[e];
            dfs(v,u);
            mxdep[u]=max(mxdep[u],mxdep[v]);
        }
        e=nxt[e];
    }
}

void getmx(int u,int fa)
{
    tmp[dep[u]]=max(tmp[dep[u]],dis[u]);
    int e=head[u];
    while(e)
    {
        int v=vet[e];
        if(v!=fa&&!flag[v]) getmx(v,u);
        e=nxt[e];
    }
}

void work(int u)
{
    flag[u]=1; 
    f[0]=0; 
    if(a[u]) K--;
    int m=0;
    int e=head[u];
    while(e)
    {
        int v=vet[e];
        if(!flag[v])
        {
            top=0;
            dep[v]=a[v];
            dis[v]=len[e];
            dfs(v,u);
            b[++m].x=mxdep[v];
            b[m].id=v;
        }
        e=nxt[e];
    }
    

    sort(b+1,b+m+1,cmp);
    for(int i=1;i<=m;i++)
    {
        int v=b[i].id;
        getmx(v,u);
        int now=0;
        if(i>1)
        {
            for(int j=b[i].x;j>=0;j--)
            {
                while(now+j<K&&now<b[i-1].x)
                {
                    now++;
                    mx[now]=max(mx[now],mx[now-1]);
                }
                if(now+j<=K) ans=max(ans,mx[now]+tmp[j]); 
            }
        }
        if(i<m)
        {
            for(int j=0;j<=b[i].x;j++)
            {
                mx[j]=max(mx[j],tmp[j]);
                tmp[j]=0;
            }
        }
         else
         {
              for(int j=0;j<=b[i].x;j++)
              {
                  if(j<=K) ans=max(ans,max(tmp[j],mx[j]));
                  tmp[j]=mx[j]=0;
             }
         }
    }
            
            
        
    if(a[u]) K++;
    e=head[u];
    while(e)
    {
        int v=vet[e];
        if(!flag[v])
        {
            sum=son[v]; root=0;
            getroot(v,u);
            work(root);
        }
        e=nxt[e];
    }

    
}

int main()
{
    freopen("spoj1825.in","r",stdin);
    freopen("spoj1825.out","w",stdout);
    int m;
    scanf("%d%d%d",&n,&K,&m);
    for(int i=1;i<=m;i++)
    {
        int x;
        scanf("%d",&x);
        a[x]=1;
    }
    for(int i=1;i<=n-1;i++)
    {
        int x,y,z;
        scanf("%d%d%d",&x,&y,&z);
        add(x,y,z);
        add(y,x,z);
    }
    ans=0;
    root=0; sum=n; c[0]=n;

    getroot(1,0);
    work(root); 
    printf("%d
",ans);
    return 0;
}