Water flowing swiftly over farmland – The August 2016 Community Challenge

This is my attempt to solve the August 2016 Community Challenge in Swift. I tried to implement the algorithm
described by @200_success:

• Each Cell keeps track of which Basin it belongs to; each Cell is initially assume to be in its own Basin. Each Basin has a

sink, or lowest Cell, which acts as a "representative element" of
the Basin, as well as a member count. Topography keeps track of
all Basins.

• For each Basin, find lowest of the sink's neighbours. If the lowest is not already a member of this Basin, transfer its cells

into the lowest neighbour's Basin, and notify Topography that the
higher basin no longer exists.

• Repeat step 2 until no further action is necessary.

• Have Topography enumerate the Basins and their counts.

The code is Swift 3 and requires Xcode 8 beta 4.

Cell.swift

import Swift

class Cell {
    let elevation: Int
    var lowerCell: Cell?
    weak var basin: Basin?

    init(elevation: Int) {
        self.elevation = elevation
    }

    /// Compare the elevation of the (currently) lowest neighbor with that of
    /// the other cell, and update `lowerCell` if necessary. Returns `true`
    /// if other cell was lower, and `false` otherwise.
    @discardableResult
    func updateLowerCell(with other: Cell) -> Bool {
        if other.elevation < (lowerCell?.elevation) ?? elevation {
            lowerCell = other
            return true
        }
        return false
    }
}

Basin.swift

```
import Swift

/// A basin consists of a "sink" and zero or more cells which eventually flow
/// into the sink.
class Basin {
var cells: [Cell]
var sink: Cell

init(cell: Cell) {
cells = [cell]
sink = cell
cell.basin = self
}

func join(with other: Basin) {
precondition(other !== self, "Basin must not be joined with itself")
for cell in other.cells {
cell.basin = self
}
cells.append(contentsOf: other

I think this code is great. Using weak references to refer to a Basin is a great way to prevent retain cycles with automatic reference counting. One suggestion I have is to make Topography a class, because it refers to a class. Other suggestions are in the code, marked with "CR". You can use this snippet of code:
/// A topography is a collection of basins. The union of all the
/// basins cells is always the complete grid.
class Topography {
var basins: [Basin]

/// Create topography from a sequence of cells. Initially, each basin
/// consists of a single cell.
init(cells: S)
where S.Element == Cell { // CR: S.Iterator.Element is the same as S.Element
basins = cells.map(Basin.init)
}

/// Create topography from 2D elevation map.
convenience init(elevationData: [[Int]]) { // CR: This calls the designated initializer, so it should be a convenience initializer.
// Create [[Cell]] array from elevationData.
let grid = elevationData.map { $0.map(Cell.init) } // CR: levelRow in clutters code, implicit names are better for simple lambdas.

// Find the lowest neighbor for each cell: left/right ...
for gridRow in grid {
for (leftCell, rightCell) in zip(gridRow, gridRow.dropFirst()) {
if !leftCell.updateLowerCell(with: rightCell) {
rightCell.updateLowerCell(with: leftCell)
}
}
}
// ... and top/down.
for (upperRow, lowerRow) in zip(grid, grid.dropFirst()) {
for (upperCell, lowerCell) in zip(upperRow, lowerRow) {
if !upperCell.updateLowerCell(with: lowerCell) {
lowerCell.updateLowerCell(with: upperCell)
}
}
}

self.init(cells: grid.flatten())
}

/// Join each basin with the next one that it is flowing into.
func joinBasins() -> Bool { // CR: classes are mutable by default, also, included return type
basins = basins.filter { basin in
if let otherBasin = basin.sink.lowerCell?.basin, otherBasin !== basin {
otherBasin.join(with: basin)
return false
}
return true // CR: Removed else, return exits
}
}

/// Sizes of all basins, sorted in decreasing order.
func basinSizes() -> [Int] {
return basins.map { $0.cells.count }.sorted(by: >)
}
}
`


StackExchange Code Review Q#137844, answer score: 4