Inspired by recent questions about counting the rooms in a floor plan (1, 2, 3), here is my attempt to solve the problem with a Swift program.

The problem (from “Counting Rooms” on CSES) is:

You are given a map of a building, and your task is to count the number of rooms. The size of the map is \$ n \times m \$ squares, and each square is either floor or wall. You can walk left, right, up, and down through the floor squares.

##Input

The first input line has two integers \$n\$ and \$m\$: the height and width of the map.

Then there are \$n\$ lines of \$m\$ characters that describe the map. Each character is . (floor) or # (wall).

##Output

Print one integer: the number of rooms.

##Constraints

\$1\le n,m \le 2500\$

Example

###Input:

5 8
########
#..#...#
####.#.#
#..#...#
########

###Output:

3

I took this as an opportunity to learn about “disjoint-set data structures” (also called “union-find data structures”). Here is my implementation, it follows closely the pseudo-code from the Wikipedia page. The only difference is that the parent of the “representative” (or “root”) does not point to itself, but is nil.

UnionFind.swift

class UnionFind {
    
    var parent: UnionFind?
    var rank: Int
    
    init() {
        self.parent = nil
        self.rank = 0
    }
    
    // Find with path compression
    func find() -> UnionFind {
        if let parent = self.parent {
            let p = parent.find()
            self.parent = p
            return p
        }
        return self
    }
    
    // Combine the trees if the roots are distinct.
    // Returns `true` if the trees were combined, and
    // `false` if the trees already had the same root.
    @discardableResult func union(with other: UnionFind) -> Bool {
        var xRoot = self.find()
        var yRoot = other.find()
        
        if xRoot === yRoot {
            return false
        }
        
        if xRoot.rank < yRoot.rank {
            swap(&xRoot, &yRoot)
        }
        
        // Merge yRoot into xRoot
        yRoot.parent = xRoot
        if xRoot.rank == yRoot.rank {
            xRoot.rank += 1
        }
        return true
    }
}

The main code is highly inspired by Efficiently counting rooms from a floorplan (version 2):

• The data is read from standard input, and can be more than one floor plan.
• For each column we keep track of the room to which the field at the current row and this column belongs to. Here we use the UnionFind structure.
• If a field is a continuation from both left and top then the two tracker values are combined with union().

main.swift

import Foundation

func readDimensions() -> (height: Int, width: Int)? {
    guard let line = readLine() else { return nil }
    let ints = line.split(separator: " ")
    guard ints.count == 2,
        let height = Int(ints[0]),
        let width = Int(ints[1]) else {
            return nil
    }
    return (height, width)
}

func readMapAndCountRooms(height: Int, width: Int) -> Int? {
    var tracker = Array<UnionFind?>(repeating: nil, count: width)
    var roomCount = 0
    
    for _ in 0..<height {
        guard let line = readLine(), line.count == width else { return nil }
        for (offset, field) in line.enumerated() {
            if field == "." {
                if let top = tracker[offset] {
                    // Continuation from the top ...
                    if offset > 0, let left = tracker[offset - 1] {
                        // ... and from the left:
                        if top.union(with: left) {
                            roomCount -= 1
                        }
                    }
                } else if offset > 0, let left = tracker[offset - 1] {
                    // Continuation from the left:
                    tracker[offset] = left
                } else {
                    // New room:
                    tracker[offset] = UnionFind()
                    roomCount += 1
                }
            } else {
                // A wall:
                tracker[offset] = nil
            }
        }
    }
    return roomCount
}

while let (height, width) = readDimensions(),
    let count = readMapAndCountRooms(height: height, width: width) {
        print(count)
}

For the input data

4 6
######
#.#..#
#....#
######
5 8
########
######.#
#......#
##.#...#
########
5 8
########
#..#...#
####.#.#
#..#...#
########
9 25
#########################
#.#.#.#.#.#.#.#.#.#.#.#.#
#########################
#.#.#.#.#.#.#.#.#.#.#.#.#
#########################
#.#.#.#.#.#.#.#.#.#.#.#.#
#########################
#.#.#.#.#.#.#.#.#.#.#...#
#########################
3 3
...
...
...
3 3
###
...
###
3 3
###
###
###
7 9
#########
#.#.#.#.#
#.#.#.#.#
#.#...#.#
#.#####.#
#.......#
#########
5 8
########
#..#.#.#
##.#.#.#
#..#...#
########
7 8
########
#..#.#.#
##.#.#.#
#..#...#
########
#..#...#
########
7 9
#########
#.#.#.#.#
#.#.#.#.#
#.#.#.#.#
#.#.#.#.#
#.......#
#########
7 9
#########
#.#.##..#
#.#.##.##
#.#.##..#
#.#...#.#
#...#...#
#########
7 9
#########
#.#.....#
#.#.###.#
#.#...#.#
#.#####.#
#.......#
#########
7 9
#########
#.......#
#.#####.#
#.#.#.#.#
#.#.#.#.#
#.......#
#########

this produces the (correct) output

1
1
3
47
1
1
0
2
2
4
1
1
1
1

Any feedback on the code is welcome, in particular:

• This is my first experiment with union-find data structures. Is it correctly implemented? Can it be simplified?
• Is the main program logic correct or can you spot any errors?