Water flowing swiftly over farmland – The August 2016 Community Challenge

Question

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

1. 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.
2. 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.
3. Repeat step 2 until no further action is necessary.
4. 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.cells)
        other.cells.removeAll()
    }
}

Topography.swift

import Swift

/// A topography is a collection of basins. The union of all the
/// basins cells is always the complete grid.
struct Topography {
    var basins: [Basin]

    /// Create topography from a sequence of cells. Initially, each basin
    /// consists of a single cell.
    init<S: Sequence>(cells: S)
        where S.Iterator.Element == Cell {
            basins = cells.map(Basin.init)
    }

    /// Create topography from 2D elevation map.
    init(elevationData: [[Int]]) {
        // Create `[[Cell]]` array from `elevationData`.
        let grid = elevationData.map { levelRow in
            levelRow.map(Cell.init)
        }

        // 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.
    mutating func joinBasins() {
        basins = basins.filter { basin in
            if let otherBasin = basin.sink.lowerCell?.basin, otherBasin !== basin {
                otherBasin.join(with: basin)
                return false
            } else {
                return true
            }
        }
    }

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

RainfallSolver.swift

import Swift

/// Wrapper type with a single static method to solve the Rainfall problem.
struct RainfallSolver {

    static func solve(elevationData: [[Int]]) -> [Int] {
        var topo = Topography(elevationData: elevationData)
        topo.joinBasins()
        return topo.basinSizes()
    }
}

The main program in main.swift which reads the elevation data (as in the specified input format) from standard input and prints the solution to standard output.

import Foundation

func readIntegers() -> [Int] {
    let line = readLine(strippingNewline: true)!
    let comps = line.components(separatedBy: CharacterSet.whitespaces).filter { !$0.isEmpty }
    return comps.map { Int($0)! }
}

func readElevationData() -> [[Int]] {
    let dimension = readIntegers().first!
    let elevationData = (0..<dimension).map { _ in readIntegers() }
    return elevationData
}

let elevationData = readElevationData()
let solution = RainfallSolver.solve(elevationData: elevationData)
// Print solution as space-separated list:
print(solution.map(String.init).joined(separator: " "))

Finally, unit tests with the four given examples from Rainfall challenge in RainfallTests.swift. All tests pass.

import XCTest

class RainfallTests: XCTestCase {

    func testExample1() {
        let solution = RainfallSolver.solve(elevationData: [
            [ 1, 5, 2 ],
            [ 2, 4, 7 ],
            [ 3, 6, 9 ]
            ])
        XCTAssertEqual(solution, [ 7, 2 ])
    }

    func testExample2() {
        let solution = RainfallSolver.solve(elevationData: [
            [ 1 ]
            ])
        XCTAssertEqual(solution, [ 1 ])
    }

    func testExample3() {
        let solution = RainfallSolver.solve(elevationData: [
            [ 1, 0, 2, 5, 8 ],
            [ 2, 3, 4, 7, 9 ],
            [ 3, 5, 7, 8, 9 ],
            [ 1, 2, 5, 4, 3 ],
            [ 3, 3, 5, 2, 1 ]
            ])
        XCTAssertEqual(solution, [ 11, 7, 7 ])
    }

    func testExample4() {
        let solution = RainfallSolver.solve(elevationData: [
            [ 0, 2, 1, 3 ],
            [ 2, 1, 0, 4 ],
            [ 3, 3, 3, 3 ],
            [ 5, 5, 2, 1 ]
            ])
        XCTAssertEqual(solution, [ 7, 5, 4 ])
    }
}

Remarks:

• Cell and Basin are defined as classes because there are references between instances of these types. In particular, the reference from a cell to its containing basin is a weak reference to avoid retain cycles. All other types are structs (value types).

• A cell does not know its coordinates. The cells are stored in a 2D grid only temporarily in init(elevationData: [[Int]]), in order to determine the lowest neighbours. After that, the actual geometry is not used anymore.

• In joinBasins(), each of the initial (single-cell) basins is visited exactly once, and potentially joined with a "lower basin" as described in step #2 of the algorithm.

  In the first version of this program, I repeated step #2 until no more basins are joined. However, after observing that the second pass never joined additional basins, I convinced myself that a single pass is sufficient. This was verified with the given examples and additional randomly created grids.

• readElevationData() in main.swift expects the data in the specified input format and will crash otherwise. However, an arbitrary amount of whitespace is accepted as separator between the integers and not just a single space character, this turned out to be convenient for my test cases. A detailed error reporting is omitted.

• The code actually works with rectangular maps, not only with quadratic maps.