Let the cows be nodes of a graph, where an edge is between every pair of cows with weight equal to the squared distance between them. Our goal is to find a Minimum Spanning Tree of this graph.
Partial Credit: Because there are a total of $\mathcal O(N^2)$ edges, we can run Kruskal's or Prim's to solve this task in $\mathcal O(N^2 \log{N})$ or $\mathcal O(N^2)$. A few modifications to the code from the Moocast analysis suffice.
Nick Wu's code:
import java.io.*;
import java.util.*;
public class MooNetworkSlow {
public static void main(String[] args) throws IOException {
BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
PrintWriter pw = new PrintWriter(System.out);
int n = Integer.parseInt(br.readLine());
int[] x = new int[n];
int[] y = new int[n];
for(int i = 0; i < n; i++) {
StringTokenizer st = new StringTokenizer(br.readLine());
x[i] = Integer.parseInt(st.nextToken());
y[i] = Integer.parseInt(st.nextToken());
}
parent = new int[n];
ArrayList<Edge> edges = new ArrayList<Edge>();
for(int i = 0; i < n; i++) {
parent[i] = i;
for(int j = 0; j < i; j++) {
long distance = (long)(x[i] - x[j]) * (x[i] - x[j]) + (long)(y[i] - y[j]) * (y[i] - y[j]);
edges.add(new Edge(i, j, distance));
}
}
Collections.sort(edges);
long sumWeights = 0;
int numComponents = n;
for(Edge curr: edges) {
if(find(curr.i) != find(curr.j)) {
merge(curr.i, curr.j);
sumWeights += curr.w;
if(--numComponents == 1) {
break;
}
}
}
pw.println(sumWeights);
pw.close();
}
static int[] parent;
public static int find(int curr) {
return parent[curr] == curr ? curr : (parent[curr] = find(parent[curr]));
}
public static void merge(int x, int y) {
parent[find(x)] = find(y);
}
static class Edge implements Comparable<Edge> {
public int i, j;
public long w;
public Edge(int a, int b, long c){
i=a;
j=b;
w=c;
}
public int compareTo(Edge e) {
return Long.signum(w-e.w);
}
}
}
Full Credit: The main idea is to reduce the number of edges we need to consider to $\mathcal O(N)$, using the special condition that $0 \leq y_i \leq 10$.
We use this important fact: consider $3$ nodes of a graph $a$, $b$, $c$. Suppose that the weights of the three edges between these nodes satisfies $w_{bc} < w_{ab}$ and $ w_{ac} < w_{ab}$. Then, the minimum spanning tree can be chosen such that edge $ab$ is not present. This can be proven by analyzing Kruskal's algorithm: the algorithm first encounters $bc$ and $ac$ in the sort order, immediately after which $a$ and $b$ must be connected. So, when $ab$ is encountered, $a$ and $b$ are already connected.
Now, given points $a = (x_1,y_1),c = (x_2,y_2), b = (x_3,y_3)$ such that $x_1 \leq x_2 < x_3$ and $y_2=y_3,$ we don't need to consider the edge between $(x_1,y_1)$ and $(x_3,y_3)$ because of the reason above. Specifically, we can easily show that $w_{bc} < w_{ab}$ and $w_{ac} < w_{ab}$.
Now, suppose that for each point $(x_1,y_1)$ we want to generate all the edges to points $(x_2,y_2)$ where $x_1 \leq x_2$ which could possibly be part of a MST. Consider the points in order of decreasing $x.$ By the above observation, we only need to consider the edges from $(x_1,y_1)$ to the last added point for each $0 \le y\le 10$, which can be maintained in linear time using an array of size $11$. Thus, we only need to run Kruskal on $11N$ distinct edges.
The overall time complexity of running this algorithm is $\mathcal O((N\max{y_i})(\log{N}+\log{\max{y_i}}))$, which comes from sorting the edge weights during Kruskal.
Danny Mittal's code:
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.*;
public class MooNetwork {
static int[] union;
static int find(int u) {
if (union[union[u]] != union[u]) {
union[u] = find(union[u]);
}
return union[u];
}
public static void main(String[] args) throws IOException {
BufferedReader in = new BufferedReader(new InputStreamReader(System.in));
int n = Integer.parseInt(in.readLine());
TreeMap<Integer, Integer>[] rows = new TreeMap[11];
for (int y = 0; y <= 10; y++) {
rows[y] = new TreeMap<>();
}
int[] xs = new int[n];
int[] ys = new int[n];
for (int j = 0; j < n; j++) {
StringTokenizer tokenizer = new StringTokenizer(in.readLine());
xs[j] = Integer.parseInt(tokenizer.nextToken());
ys[j] = Integer.parseInt(tokenizer.nextToken());
rows[ys[j]].put(xs[j], j);
}
List<Edge> edges = new ArrayList<>();
for (int j = 0; j < n; j++) {
for (int y = 0; y <= 10; y++) {
Map.Entry<Integer, Integer> entry = y == ys[j] ? rows[y].lowerEntry(xs[j]) : rows[y].floorEntry(xs[j]);
if (entry != null) {
long dx = xs[j] - entry.getKey();
long dy = ys[j] - y;
edges.add(new Edge(j, entry.getValue(), (dx * dx) + (dy * dy)));
}
}
}
union = new int[n];
for (int j = 0; j < n; j++) {
union[j] = j;
}
edges.sort(Comparator.comparingLong(edge -> edge.weight));
long answer = 0;
for (Edge edge : edges) {
int u = find(edge.a);
int v = find(edge.b);
if (v != u) {
union[v] = u;
answer += edge.weight;
}
}
System.out.println(answer);
}
static class Edge {
final int a;
final int b;
final long weight;
Edge(int a, int b, long weight) {
this.a = a;
this.b = b;
this.weight = weight;
}
@Override
public String toString() {
return "Edge{" +
"a=" + a +
", b=" + b +
", weight=" + weight +
'}';
}
}
}
Fun Fact #1: Note that Squared Euclidean Distance was not very special in our proof. In particular, the same solution can be used for Manhattan Distance or regular Euclidean Distance.
Fun Fact #2: Note that finding an MST in Squared Euclidean Distance is equivalent to finding an MST in Euclidean Distance, although it is not equivalent to finding an MST in Manhattan distance (can you see why?). So, an algorithm that can find a Euclidean MST could be used to solve this problem. General Euclidean MST (without the bound on $y_i$) can be done in $\mathcal O(N \log{N})$ time.