Acwing 3728. 城市通电 java

⭐ 原题链接

⭐ prim 最小生成树 O(n^2)

⭐ 记得开 long

import java.io.*;
import java.util.*;

public class Main
{
	static BufferedReader in = new BufferedReader(new InputStreamReader(System.in));
	static BufferedWriter out = new BufferedWriter(new OutputStreamWriter(System.out));

	static int N = 2020, n;
	static int INF = 0x3f3f3f3f;


	static boolean[] st = new boolean[N];// 记录第 i 个点是否在连通块里边
	static long[] dist = new long[N];// dist 记录的是 i 到当前已知生成树的距离
	static long[] distp = new long[N];// 某点的父节点，即由哪个点扩展而来的，0 表示 是超级源点直连，即 建站
	static Node[] pos = new Node[N];// 城市坐标数组
	static int[] cp = new int[N];// city price 存在 i 城市建造一个发电站的费用
	static int[] k = new int[N];// 存每个城市的 常数 k

	static List<Integer> l1 = new ArrayList<>();// 记录建造发电站的城市
	static List<Node> l2 = new ArrayList<>();// 记录拉电线（边）的村庄

	static class Node
	{
		long x;
		long y;

		public Node(long x, long y)
		{
			super();
			this.x = x;
			this.y = y;
		}
	}

	static long prim()
	{
		long res = 0;
		Arrays.fill(dist,INF);
		st[0] = true;
		dist[0] = 0;
//		根节点是 超级源点 到任何一个城市的权重就是 该城市建站的费用
		for (int i = 1; i <= n; i++)
			dist[i] = cp[i];

		for (int i = 0; i < n; i++)//枚举 n 次把每个点都加入到最小生成树中
		{
		  //  找出距离当前点集合中距离最近的点
			int t = -1;
			for (int j = 1; j <= n; j++)//枚举每一个点，找出不在当前生成树的最短距离的点
				if (!st[j] && (t == -1 || dist[t] > dist[j]))
					t = j;

			st[t] = true;
			res += dist[t];
			
			if (distp[t] == 0)  //建站
				l1.add(t);
			else//普通村庄
			{
				l2.add(new Node(distp[t], t));
			}
//          用新加入的点更新其他点的距离
			for (int j = 1; j <= n; j++)
			{
				if (!st[j] && dist[j] > getEdge(t, j))
				{
					dist[j] = getEdge(t, j);// dist 记录的是 j 到当前生成树的距离
					distp[j] = t;//记录 j 的父节点 t
				}
			}

		}
		return res;
	}

//	求 a b 两个城市中边的长度
	static long getEdge(int a, int b)
	{
		long kk = k[a] + k[b];
		long dx = pos[a].x - pos[b].x;
		long dy = pos[a].y - pos[b].y;

		return kk * (Math.abs(dx) + Math.abs(dy));
	}

	public static void main(String[] args) throws IOException
	{
		n = Integer.parseInt(in.readLine());

//		输入城市位置
		for (int i = 1; i <= n; i++)
		{
			String[] split = in.readLine().split(" ");
			int x = Integer.parseInt(split[0]);
			int y = Integer.parseInt(split[1]);
			pos[i] = new Node(x, y);
		}

//		城市建站的费用
		String[] split = in.readLine().split(" ");
		for (int i = 1; i <= n; i++)
		{
			int x = Integer.parseInt(split[i - 1]);
			cp[i] = x;
		}

//		城市的电线单位价格
		String[] split2 = in.readLine().split(" ");
		for (int i = 1; i <= n; i++)
		{
			k[i] = Integer.parseInt(split2[i - 1]);
		}

		System.out.println(prim());
		System.out.println(l1.size());
		for (int x : l1)
		{
			System.out.print(x + " ");
		}
		System.out.println();
		System.out.println(l2.size());
		for (Node n : l2)
		{
			System.out.println(n.x + " " + n.y);
		}

	}

}