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 whichBasin
it belongs to; eachCell
is initially assume to be in its ownBasin
. EachBasin
has asink
, or lowestCell
, which acts as a "representative element" of theBasin
, as well as a member count.Topography
keeps track of allBasin
s.- For each
Basin
, find lowest of the sink's neighbours. If the lowest is not already a member of thisBasin
, transfer its cells into the lowest neighbour'sBasin
, and notifyTopography
that the higher basin no longer exists.- Repeat step 2 until no further action is necessary.
- Have
Topography
enumerate theBasin
s 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
andBasin
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 ininit(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.