[BZOJ 3872][POI 2014]Ant colony(BFS+二分)

题目链接

http://www.lydsy.com/JudgeOnline/problem.php?id=3872

思路

很显然可以考虑倒着来搞，让蚂蚁从食蚁兽所在的两个点开始，倒着往回走，就能得到每个点上至少和至多有多少个蚂蚁，才能让食蚁兽所在的点刚好有 k 个蚂蚁。这样就能在每个点得到一个区间。假如 x 个蚂蚁这次在点 t ,上一次在点 t′ ，则可以得到上一次蚂蚁的个数 x′ 区间为

[max{degree[t]−1,1}x,max{degree[t]−1,1}(x+1)−1]

这样就能得到在每个叶节点，要让恰好 k 只蚂蚁到食蚁兽所在的点所需的蚂蚁个数的区间，然后用lower_bound和upper_bound来二分出 g 群蚂蚁里哪些蚂蚁群的个数是在这个区间里的。

代码

#include <iostream>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <algorithm>

#define MAXE 1000010
#define MAXV 1000010

using namespace std;

typedef long long int LL;

int S,T;
int m[MAXV];
LL INF;

struct edge
{
    int u,v,next;
}edges[MAXE*2];

int head[MAXV],nCount=0;

void AddEdge(int U,int V)
{
    edges[++nCount].u=U;
    edges[nCount].v=V;
    edges[nCount].next=head[U];
    head[U]=nCount;
}

int degree[MAXV];
int n,g,K;
int q[MAXE*4];
LL L[MAXV],R[MAXV],ans=0;
bool vis[MAXV];

void BFS()
{
    memset(vis,false,sizeof(vis));
    int h=0,t=0;
    L[S]=(LL)K*(LL)max((degree[S]-1),1);
    R[S]=min((LL)(K+1)*(LL)max((degree[S]-1LL),1LL)-1LL,(LL)INF);
    if(L[S]<INF) q[t++]=S;
    L[T]=(LL)K*(LL)max((degree[T]-1LL),1LL);
    R[T]=min((LL)(K+1LL)*(LL)max((degree[T]-1LL),1LL)-1LL,(LL)INF);
    if(L[T]<INF) q[t++]=T;
    vis[S]=vis[T]=true;
    while(h<t)
    {
        int u=q[h++];
        if(degree[u]==1)
        {
            int r=upper_bound(m+1,m+g+1,R[u])-m,l=lower_bound(m+1,m+g+1,L[u])-m;
            ans+=K*(r-l);
        }
        //else
        {
            for(int p=head[u];p!=-1;p=edges[p].next)
            {
                int v=edges[p].v;
                if(vis[v]) continue;
                vis[v]=true;
                L[v]=L[u]*(LL)(max(degree[v]-1LL,1LL));
                R[v]=min((R[u]+1LL)*(LL)(max(degree[v]-1LL,1LL))-1LL,(LL)INF);
                if(L[v]<INF) q[t++]=v;
            }
        }
    }
}

int main()
{
    memset(head,-1,sizeof(head));
    scanf("%d%d%d",&n,&g,&K);
    for(int i=1;i<=g;i++) scanf("%d",&m[i]);
    sort(m+1,m+g+1);
    INF=m[g]+1;
    scanf("%d%d",&S,&T);
    AddEdge(S,T);
    AddEdge(T,S);
    degree[S]++,degree[T]++;
    for(int i=2;i<n;i++)
    {
        int u,v;
        scanf("%d%d",&u,&v);
        AddEdge(u,v);
        AddEdge(v,u);
        degree[u]++,degree[v]++;
    }
    BFS();
    printf("%lld\n",ans);
    return 0;
}