Problem overview
The function minCostConnectPoints receives a slice points, where each point is an [x, y] pair. Connecting two points costs their Manhattan distance, |x1 - x2| + |y1 - y2|: the distance along the grid lines. The function returns the least total cost that connects all the points, meaning you can get from any point to any other by following connections.
For [[0,0],[2,2],[3,10],[5,2],[7,0]] the answer is 20, from four connections: [0,0]-[2,2] costs 4, [2,2]-[5,2] costs 3, [5,2]-[7,0] costs 4 and [2,2]-[3,10] costs 9. For [[3,12],[-2,5],[-4,1]] the answer is 18, and for a single point it is 0, because there is nothing to connect.
Approach: a minimum spanning tree with Kruskal's algorithm
Think of each pair of points as an edge of a graph, with the Manhattan distance as its cost. There are n · (n - 1) / 2 edges. The cheapest way to connect everything is a minimum spanning tree: n - 1 edges that join all points with the least total cost. A tree has no cycles, and any extra edge would only add cost.
Kruskal's algorithm builds it greedily. Sort all edges from cheapest to dearest. Go through them in that order and take an edge whenever its two ends are in different groups, which merges the two groups. If both ends are already in the same group, the edge would close a cycle, so skip it. After n - 1 edges everything is in one group, and the loop can stop.
Taking the cheapest edge between two groups is always safe: the finished tree has to join those two groups somehow, and swapping a dearer connection for the cheaper one would give a tree that costs no more.
To ask "are these two points in the same group?" quickly, use union-find. Each group is a tree of points. Every point links to a parent, and the root, which is its own parent, names the group. Two points are in the same group exactly when they lead to the same root. With union by size and path halving, one lookup costs almost constant time.
Example 3, [[0,0],[1,0],[0,1],[10,10]], shows the idea. The six edge costs sorted are 1, 1, 2, 19, 19 and 20. The two edges of cost 1 join [0,0] with [1,0] and with [0,1]. The edge of cost 2, between [1,0] and [0,1], joins two points that are already connected through [0,0], so it is skipped. The next edge of cost 19 attaches [10,10]. That is three edges, n - 1, and the total is 1 + 1 + 19 = 21.
Prim's algorithm is the other classic way to build a minimum spanning tree. It grows one tree from a starting point, but this lesson uses Kruskal's.
Step by step
- Store the point count as
n := len(points). - Create
edgesas[][]int{}. Each edge will be[]int{cost, i, j}, with the cost first because that is the value the edges are sorted by. - Loop over
ifrom 0 ton - 1. - Loop over
jfromi + 1ton - 1. This lists each pair once and never pairs a point with itself, because connectingitojcosts the same asjtoi. - Compute
costasabsof the x difference plusabsof the y difference. The absolute value makes each difference a distance, whichever point is further left or lower. The lesson provides an integerabs, sincemath.Absonly works onfloat64. - Append the edge and store the result back in
edges. - Sort
edgesby cost withsort.Slice, passing a function that says whetheredges[i][0] < edges[j][0]. - Create
parentas a slice ofnints.parent[x]will hold the parent of pointx. - Create
sizethe same way.size[r]counts the points in the group whose root isr. - Loop over every point to initialise both slices.
- Set
parent[i] = i. A point that is its own parent is a root, so at the start every point is a group of one. - Set
size[i] = 1.makefilled the sizes with zeros, so they must be set. - Declare
total := 0, the cost of the edges taken so far. - Declare
used := 0, the number of edges taken. A tree overnpoints has exactlyn - 1. - Loop over
edgeswithrange, cheapest first.edge[0]is the cost, andedge[1]andedge[2]are the two points. - Find the root of
edge[1]withfind(parent, edge[1])and keep it inrootA. The helper is written at the end. - Find the root of
edge[2]inrootB. - If
rootA != rootB, the edge joins two groups, and the network takes it. Equal roots mean the edge would close a cycle. - Make
rootAthe root of the larger group. Ifsize[rootA] < size[rootB], swap the two withrootA, rootB = rootB, rootA. This is union by size: the smaller tree goes under the larger one, which keeps every path short. - Set
parent[rootB] = rootAto merge the groups. - Add
size[rootB]tosize[rootA]. The merged group has the points of both. - Add
edge[0]tototal. This is the key line: the edge is the cheapest one left that joins two groups. - Add 1 to
used. - Return
totalwhenused == n - 1. Every point is in one group, and the remaining edges could only close cycles. - Return
totalafter the loop. That is reached only with a single point: there are no edges, and the cost is 0. With two or more points the function always returns from inside the loop, since every pair is an edge. - Write
find(parent []int, x int) intbelow the main function. It returns the root ofx's group. - Loop while
parent[x] != x. A Goforwith only a condition works like awhile, and a loop needs no call stack however long the walk is. - Point
xto its grandparent:parent[x] = parent[parent[x]]. This is path halving, and later walks from here take about half the steps. - Move up with
x = parent[x]. - Return
x, the root.
The solution in Go
func minCostConnectPoints(points [][]int) int {
n := len(points)
edges := [][]int{}
for i := 0; i < n; i++ {
for j := i + 1; j < n; j++ {
cost := abs(points[i][0] - points[j][0]) + abs(points[i][1] - points[j][1])
edges = append(edges, []int{cost, i, j})
}
}
sort.Slice(edges, func(i, j int) bool {
return edges[i][0] < edges[j][0]
})
parent := make([]int, n)
size := make([]int, n)
for i := 0; i < n; i++ {
parent[i] = i
size[i] = 1
}
total := 0
used := 0
for _, edge := range edges {
rootA := find(parent, edge[1])
rootB := find(parent, edge[2])
if rootA != rootB {
if size[rootA] < size[rootB] {
rootA, rootB = rootB, rootA
}
parent[rootB] = rootA
size[rootA] += size[rootB]
total += edge[0]
used++
if used == n - 1 {
return total
}
}
}
return total
}
func find(parent []int, x int) int {
for parent[x] != x {
parent[x] = parent[parent[x]]
x = parent[x]
}
return x
}
The Go details worth a closer look:
- Edges as
[]int. Each edge is a slice of three ints. A struct would be clearer in a larger program, but a slice keeps the sort code short here. sort.Slice. It takes the slice and a function that receives two positions and reports whether the element at the first should come before the one at the second. It is not stable: edges of equal cost can end up in either order. That is fine, because every minimum spanning tree has the same total cost.- Swapping in one line.
rootA, rootB = rootB, rootAevaluates the whole right side before it assigns, so no temporary variable is needed. - Slices share their memory.
findreceivesparentas a slice, which points to the same underlying array. Whenfindhalves a path,minCostConnectPointssees the change. Assigning to the parameterx, in contrast, changes onlyfind's own copy of the number. - Iteration, not recursion.
finduses a loop, so a long chain of parents can't overflow the call stack. intis 64 bits. Coordinates reach 1,000,000 in size, so one cost is at most 4,000,000 and the total stays far below the limit.
Complexity
Time: O(n² log n). Listing the edges takes O(n²), and sorting n · (n - 1) / 2 of them takes O(n² log n). The loop makes two find calls per edge, each close to constant time with union by size and path halving, so the loop is O(n²) up to that tiny factor and the sort dominates. For 1000 points there are 499,500 edges.
Space: O(n²). The edge list holds every pair. parent and size need only O(n).
Edge cases and common mistakes
- A single point.
[[5, 5]]returns 0: there are no edges, and the loop is skipped. - Two points.
[[0, 0], [3, 4]]returns 7, the Manhattan distance, not the straight-line 5. - Far-apart points.
[[-1000000, -1000000], [1000000, 1000000]]returns 4000000, which fits easily in anint. - Points on a line.
[[0,0],[0,1],[0,2],[0,3]]returns 3: the tree chains neighbours. - Using
math.Abs. It takes and returnsfloat64, so using it on ints needs conversions in both directions. A small integerabsavoids that. - Taking edges without the group check. Some cheap edges close a cycle. In example 3 the edge of cost 2 joins two points that are already connected, and taking it would only add cost to the total.
- Counting an edge twice. Starting
jat 0 instead ofi + 1adds self-edges of cost 0 and every pair twice. The self-edges are harmless to the answer, because both ends share a root, but they double the list and the sorting work. - Stopping at
nedges. A tree hasn - 1. Checkingused == n - 1stops the loop at the right time.
Testing
[[0,0],[2,2],[3,10],[5,2],[7,0]]gives20.[[3,12],[-2,5],[-4,1]]gives18.[[0,0],[1,0],[0,1],[10,10]]gives21.[[5,5]]gives0.[[0,0],[3,4]]gives7.[[-1000000,-1000000],[1000000,1000000]]gives4000000.[[0,0],[0,1],[0,2],[0,3]]gives3.
