Problem overview
The function minCostConnectPoints(_ points: [[Int]]) -> Int receives n points, each an [x, y] pair. Connecting two points costs their Manhattan distance, |x1 - x2| + |y1 - y2|, the distance along grid lines. The function returns the least total cost that makes the points all connected, 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: the connections [0, 0]-[2, 2] (4), [2, 2]-[5, 2] (3), [5, 2]-[7, 0] (4) and [2, 2]-[3, 10] (9) add up to 20. For [[3, 12], [-2, 5], [-4, 1]] the answer is 18, and a single point costs 0. There are up to 1000 points, with coordinates between -1,000,000 and 1,000,000.
Approach: Kruskal's algorithm
Treat every pair of points as an edge of a graph whose cost is the pair's Manhattan distance. There are n · (n - 1) / 2 edges. The cheapest way to connect all the points is a minimum spanning tree: n - 1 edges that connect every point with the smallest total cost.
Kruskal's algorithm builds it greedily. Sort all edges from cheapest to dearest. Go through them in order and take an edge when its two ends are still in different groups, which joins the two groups into one. If both ends are already in the same group, the edge would only close a cycle and add cost, so skip it. After n - 1 edges taken, everything is one group and the algorithm can stop.
Taking the cheapest edge between two groups is always safe. The finished tree has to join those two groups somehow, and if it used a dearer edge there, swapping in the cheaper one would give a tree that costs no more.
To know quickly whether two points are in the same group, the code uses union-find. Each group is a tree of parent links, and the root, the one point that is its own parent, names the group. Two points are in the same group exactly when they lead to the same root.
On [[0, 0], [1, 0], [0, 1], [10, 10]] the edges [cost, i, j] come out of the double loop as:
[[1, 0, 1], [1, 0, 2], [20, 0, 3], [2, 1, 2], [19, 1, 3], [19, 2, 3]]
and sorted by cost:
[[1, 0, 1], [1, 0, 2], [2, 1, 2], [19, 1, 3], [19, 2, 3], [20, 0, 3]]
The two cost-1 edges join points 0, 1 and 2. The edge of cost 2 joins two points that are already together, so it is skipped. The first edge of cost 19 brings in point 3. That is the third edge, n - 1 = 3, so the total is 1 + 1 + 19 = 21 and the loop returns.
The other classic method is Prim's algorithm, which grows one tree from a starting point. This lesson uses Kruskal's.
Step by step
- Store
n, the number of points, withlet. - Create
edgesas an empty[[Int]]. Each edge will be[cost, i, j], with the cost first because it is the sort key. - Loop
iover0..<nandjover(i + 1)..<n. Startingjati + 1lists each pair once and never pairs a point with itself, since connectingitojcosts the same asjtoi. Wheniis the last index the range is empty, which is fine. - Compute
costasabsof the x difference plusabsof the y difference, and append[cost, i, j]toedges. - Sort
edgesby cost with a closure that compares$0[0]and$1[0]. - Set up union-find:
parentisArray(0..<n), so every point starts as its own root, andsizeisnones, counting the points in each group. - Declare
totalandusedat 0: the cost so far and the number of edges taken. - Loop over
edges, cheapest first. Find the rootsrootAandrootBofedge[1]andedge[2]. They arevarbecause the next step may swap them. - If the roots differ, take the edge. Make
rootAthe root of the larger group by swapping the two whensize[rootA] < size[rootB]. This is union by size: the smaller tree goes under the larger one, which keeps the trees shallow. - Merge: set
parent[rootB]torootAand addsize[rootB]tosize[rootA]. - Add
edge[0]tototaland add 1 toused. This is the key line: the edge is the cheapest one left that joins two groups. - When
used == n - 1, returntotal. Every point is in one group. - After the loop, return
total. The loop only gets here when there is a single point, which has no edges and costs nothing. - Write
findas a method that takes_ parent: inout [Int]and a node. Copy the node into avar x, and whileparent[x] != xpointparent[x]to its grandparent and movextoparent[x]. Returnx, the root.
The solution in Swift
class Solution {
func minCostConnectPoints(_ points: [[Int]]) -> Int {
let n = points.count
var edges = [[Int]]()
for i in 0..<n {
for j in (i + 1)..<n {
let cost = abs(points[i][0] - points[j][0]) + abs(points[i][1] - points[j][1])
edges.append([cost, i, j])
}
}
edges.sort { $0[0] < $1[0] }
var parent = Array(0..<n)
var size = [Int](repeating: 1, count: n)
var total = 0
var used = 0
for edge in edges {
var rootA = find(&parent, edge[1])
var 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 += 1
if used == n - 1 {
return total
}
}
}
return total
}
func find(_ parent: inout [Int], _ node: Int) -> Int {
var x = node
while parent[x] != x {
parent[x] = parent[parent[x]]
x = parent[x]
}
return x
}
}
The Swift details worth a closer look:
inoutand&. Arrays are value types, so a function would normally get a copy. Marking the parameterinoutletsfindchange the caller'sparent, and the caller passes it with&. A copy per call would turn each of the roughly half a million edges into a full-array copy.- The sort closure.
edges.sort { $0[0] < $1[0] }sorts in place.$0and$1are the two arguments, and the closure says whether the first goes before the second. Edges with equal costs may come out in any order, which doesn't change the total. - Tuple swap.
(rootA, rootB) = (rootB, rootA)builds a pair from the right side and unpacks it into the left, so a swap needs no temporary variable. - Path halving.
parent[x] = parent[parent[x]]makes each visited node skip one level, so later walks are shorter. With union by size, a call costs almost constant time on average. The loop needs no recursion, so a long chain can't overflow the call stack. Array(0..<n). This turns a range into the array[0, 1, ..., n - 1], soparent[i]starts asi.- A half-open range can be empty.
(i + 1)..<nis empty wheniis the last index, and1..<1is empty for one point. Both are fine. A closed range(i + 1)...(n - 1)that ends before it starts would stop the program. - Integer size. A cost is at most 4,000,000, and the total at most about 4 billion. Swift's
Intis 64 bits, so neither comes near an overflow.
Complexity
Time: O(n² log n). Building the n(n - 1)/2 edges takes O(n²). Sorting them takes O(n² log n), and that dominates. Each edge then needs two find calls, which are nearly constant time, so the loop costs O(n²) up to that small factor.
Space: O(n²). The edge list holds one three-number array per pair. parent and size add only O(n).
Edge cases and common mistakes
- A single point.
[[5, 5]]gives 0. There are no edges, and the loop is skipped. - Two points.
[[0, 0], [3, 4]]gives 7, the Manhattan distance, not the straight-line one (which would be 5). - Large coordinates.
[[-1000000, -1000000], [1000000, 1000000]]gives 4000000. Subtracting values of opposite signs needs no special care in a 64-bitInt. - Taking every cheap edge. The edge of cost 2 in the example above closes a cycle through point 0. Taking it would add 2 to the total for nothing. The root comparison is what rejects it.
- Comparing
parentvalues instead of roots. The parent of a point is not its group's name until you follow the links to the root, so always callfind. - Not stopping early. The answer stays correct without the
used == n - 1check, since the remaining edges are all skipped, but it does pointless work over the rest of the sorted list. - Pairing a point with itself. Starting
jatiwould add zero-cost self-loops that are skipped but waste time and memory.
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.[[0, 0], [1, 1], [1, 0], [-1, 1]]gives4.
