牛客练习赛41——C-抓捕盗窃犯

链接：https://ac.nowcoder.com/acm/contest/373/C
来源：牛客网
 

时间限制：C/C++ 1秒，其他语言2秒
空间限制：C/C++ 32768K，其他语言65536K
64bit IO Format: %lld

题目描述

Q市发生了一起特大盗窃案。这起盗窃案是由多名盗窃犯联合实施的,你要做的就是尽可能多的抓捕盗窃犯。
已知盗窃犯分布于N个地点,以及第i个地点初始有ai名盗窃犯。
特别的是，对于每一个地点u,都有一个固定的地点v--当前如果某个盗窃犯位于地点u,在下一个时刻他会移动到地点v。
你需要通过初始时在某些点设置哨卡来捉住他们。
现在你可以在M个地点设置哨卡,如果在某个地点设置哨卡,你可以抓获在任一时刻经过该地点的盗窃犯。
也就是说,哨卡存在的时间是无限长,但哨卡不能移动。

输入描述:

第一行两个整数N,M(1≤N,M≤10^5)。
第二行N个整数a1a2...aN(0≤a1,a2,...aN≤10^5),表示第i个地点初始有ai名盗窃犯。
第三行N个整数v1v2...vN(1≤v1,v2,...vN≤N),表示当前处于地点i的盗窃犯下一个时刻会移动到地点vi。

输出描述:

输出一行一个整数--能够抓捕到的最大数量。

示例1

输入

8 2
1 2 3 4 1 2 3 12
2 3 3 3 6 7 5 8

输出

22

说明

对于样例一，一种可行的方案是：

在地点3、地点8分别设置一个哨卡，此时答案为1+2+3+4+12=22

示例2

输入

8 2
1 2 3 4 5 6 7 8
2 3 4 5 6 7 8 8

输出

36

说明

对于样例二，一种可行的方案是：

在地点2、地点8分别设置一个哨卡，此时答案为1+2+3+4+5+6+7+8=36

解题思路：

        由于每个节点都有且仅有一个后继节点，因此至少会形成一个环（自己指向自己也可以算作一个环）。所有在环外的节点的人数必定会流向环内。所以需要求出每个环的总人数，再排序按照哨卡数求和即可。在求出环的总人数之前，需要把环外的人数累加到对应的环内。这里使用拓扑排序寻找不在环内的节点并将人数累加到后继节点上，直到不存在前驱节点个数为0的节点，则剩余节点都是组成环的节点，再遍历每个环即可。

        拓扑排序需要使用栈实现，循环查找前驱节点数为0的节点会使时间复杂度大大增加。当节点近似于链表形状时，累加人数可能会超过int最大值，所以数据都使用long long存储。

AC代码：

#include <iostream>
#include <algorithm>
#include <stack>

using namespace std;

int main(){
	int n, m;
        //int a[100001]; WA
	long long a[100001];
	int v[100001];
	int count[100001] = {};
	cin >> n >> m;
	for(int i = 1; i <= n; i++){
		cin >> a[i];
	}
	for(int i = 1; i <= n; i++){
		cin >> v[i];
		//记录父节点个数 
		count[v[i]]++;
	}
        /* TLE
	    int index;
	    do{
		    index = -1;
		    for(int i = 1; i <= n; i++){
			    if(count[i] == 0){
			    	index = i;
			    	break;
			    }
		    }
		    if(index != -1){
			    a[v[index]] += a[index];
			    a[index] = 0;
			    count[index] = -1;
			    count[v[index]]--;
		    }
	    }while(index != -1);
        */
	stack<int> s;
	//寻找没有父节点的节点 
	for(int i = 1; i <= n; i++){
		if(count[i] == 0){
			s.push(i);
		}
	}
	while(!s.empty()){
		int index = s.top();
		s.pop();
		//将它的人数合并到子节点上 
		a[v[index]] += a[index];
                //让调试时显示更直观，无用
		a[index] = 0;
		count[index] = -1;
		//子节点的父节点个数-1 
		count[v[index]]--;
		if(count[v[index]] == 0){
			s.push(v[index]);
		}
	}
	//没有父节点的节点处理完毕 剩下的全是环
	long long num[100001] = {};
	int huan_count = 0;
	for(int i = 1; i <= n; i++){
		//存在环 
		if(count[i] > 0){
			int next = v[i];
			//遍历环，累加人数 
			while(next != i){
				num[huan_count] += a[next];
				//删除这个环内所有元素的引用 
				count[next] = -1;
				next = v[next];
			}
			//加上自身
			num[huan_count] += a[i];
			huan_count++;
		}
	}
        //升序排序
	sort(num, num + huan_count);
	long long result = 0;
	for(int i = 0, j = huan_count - 1; i < m && j >= 0; i++, j--){
		result += num[j];
	}
	cout << result << endl;
}

优化解法：

        （虽然说是优化不过耗时几乎与上面相同）此题观察后可以发现，每个节点仅且仅有一个后继节点，最终会构成1+个连通图。每个节点初始都可以看为一个单独的节点，通过vi相互联通，因此可以通过并查集计算各集合内的最大人数。

优化AC代码：

#include <iostream>
#include <map>
#include <algorithm>

using namespace std;

int index[100001] = {};

int find(int x){
	//如果查询的这个是根节点（并查集中节点没有主次之分，集合内任何一个节点都可以成为根节点），则返回，否则继续查询这个节点父节点是否为根节点，同时更新路径上所有节点的根节点 
	return x == index[x] ? x : index[x] = find(index[x]);
}

int main(){
	int n, m;
	int a[100001];
	int v[100001];
	cin >> n >> m;
	for(int i = 1; i <= n; i++){
		cin >> a[i];
		//初始化，使每个节点成为一个单独的集合 
		index[i] = i;
	}
	for(int i = 1; i <= n; i++){
		cin >> v[i];
	}
	for(int i = 1; i <= n; i++){
		//根据题目将相同的集合合并 
		index[find(v[i])] = find(i);
	}
	long long sum[100001] = {};
	for(int i = 1; i <= n; i++){
		//根据根节点累加（因为数组中路径不一定完全更新了，可以输入示例1调试，在此循环之前index数组前四个值为4,1,1,4，但他们应属于同一集合，只是2、3的根节点没有及时更新，再次调用find后会更新为4,4,4,4） 
		sum[find(i)] += a[i];
	} 
	sort(sum + 1, sum + n + 1);
	long long result = 0;
	for(int i = n, j = 0; i >= 0 && j < m; i--, j++){
		result += sum[i];
	}
	cout << result << endl;
}