小美的修路(最小生成树练习)

本题链接：登录—专业IT笔试面试备考平台_牛客网

题目：

样例：

输入

3 4 1 2 3 1 1 2 2 0 1 3 1 0 2 3 3 0

输出

2 1 3

思路：

        由题意，这里建造的城市需要修路，且每个城市之间可以联通，且 是 1 的标记，一定有该方案，0 可自主选择该修路方案，问最少花费修路费用，。

根据题干 ‘每个城市之间可以联通’ 相当于 每个结点都需要遍历一遍，这个修路，就是边权。

这里只是多了一个 标记需要优先选择，根据数据范围，我们还是用 Kruskal 算法即可。

代码详解如下：

#include <iostream>
#include <vector>
#include <queue>
#include <cstring>
#include <algorithm>
#include <unordered_map>
#define endl '\n'
#define YES puts("YES")
#define NO puts("NO")
#define umap unordered_map
#define All(x) x.begin(),x.end()
#pragma GCC optimize(3,"Ofast","inline")
#define IOS std::ios::sync_with_stdio(false),cin.tie(0), cout.tie(0)
using namespace std;
const int N = 2e6 + 10;

int n,m;

struct Edge
{
	int a,b,w,p,id;
	// 定义排序规则，标记 1 的为优先选择方案
	// 之后排序 最小花费的边权为 优先可选择的方案
	inline bool operator<(const Edge&t)const
	{
		if(p != t.p) return p > t.p;
		return w < t.w;
	}
}edge[N];

umap<int,int>r;	// 存储集合所对应的连接点

// 查找对应的城市 根节点函数
inline int Find(int &x)
{
	int t = x;
	while(x != r[x]) x = r[x];
	r[t] = x;
	return x;
}

vector<int>plan;
inline bool Kruskal()
{
	// 排序好优先选择的方案
	sort(edge + 1,edge + m + 1);
	
	// 初始化城市点的根节点为自身
	for(int i = 0;i <= n;++i) r[i] = i;
	
	int cnt = 0;	// cnt 用于记录修路的数量
	// 遍历所有方案
	for(int i = 1;i <= m;++i)
	{
		// 获取需要修路的对应两个城市
		int a = edge[i].a;
		int b = edge[i].b;
		
		// 查找对应两个城市的根节点
		a = Find(a),b = Find(b);
		
		if(edge[i].p || a != b)
		{
			// 如果这两个城市之间没有连接过
			// 或者它们是必选方案，那么将它们连接起来
			r[a] = b;
			// 累加方案数
			plan.emplace_back(edge[i].id);
			++cnt;	// 累加可以修的路数量
		}
	}
	// 如果所修的路无法将所有城市联通，返回 false
	if(cnt < n - 1) return false;
	return true;	// 否则返回 true
}

// 打印方案数函数
inline void PrintPlan()
{
	cout << plan.size() << endl;
	for(int i : plan)
	{
		cout << i << ' ';
	}
}

inline void solve()
{
	// 输入各个信息
	cin >> n >> m;
	for(int i = 1;i <= m;++i)
	{
		int a,b,w,p;
		cin >> a >> b >> w >> p;
		// 存储好方案
		edge[i] = {a,b,w,p,i};
	}
	
	// 开始克鲁斯卡尔算法，判断是否有解，并输出对应答案
	if(Kruskal()) PrintPlan();
	else puts("-1");
}

int main()
{
//	freopen("a.txt", "r", stdin);
	IOS;
	int _t = 1;
//	cin >> _t;
	while (_t--)
	{
		solve();
	}

	return 0;
}

最后提交：