Swift) 백준 16234번 인구이동 - G4

https://www.acmicpc.net/problem/16234

16234번: 인구 이동

dfs를 사용해서 그룹으로 묶은 후 그룹의 평균값 구해서 다시 배열에 삽입. 만약 리턴된 그룹의 count 값이 1개보다 커야지만 그룹 배열에 그룹을 추가. for문을 다 돈후 그룹 배열의 count가 0이면은 더이상 인구를 나눌 수 없으므로 break하게 했다.

날짜수 카운트는 인구가 다 이동하고 난 후 카운트가 되도록 하는 것이 문제 조건이므로 이를 유의해야 함.

import Foundation

let nlr = readLine()!.split(separator: " ").map { Int($0)! }
let n = nlr[0]
let l = nlr[1]
let r = nlr[2]

var worldMap: [[Int]] = []
var isGroupMap: [[Bool]] = Array(repeating: Array(repeating: false, count: n), count: n)
var answer = 0

for _ in 0..<n {
    let input = readLine()!.split(separator: " ").map { Int($0)! }
    worldMap.append(input)
}

func makeGroup(i: Int, j: Int, group: [(Int, Int)]) -> [(Int, Int)] {
    if isGroupMap[i][j] {
        return group
    }
    var cpGroup = group
    isGroupMap[i][j] = true
    cpGroup.append((i, j))
    let ni = [0,1,-1,0]
    let nj = [1,0,0,-1]
    for num in 0..<4 {
        let mi = i+ni[num]
        let mj = j+nj[num]
        if mi < 0 || mi >= n || mj < 0 || mj >= n {
            continue
        }
        if isGroupMap[mi][mj] {
            continue
        }
        let countryDiff = abs(worldMap[i][j] - worldMap[mi][mj])
        if countryDiff >= l && countryDiff <= r {
            cpGroup.append(contentsOf: makeGroup(i: mi, j: mj, group: []))
        }
    }
    return cpGroup
}

func calculateWorldMap(groupArr: [(Int, Int)]) {
    var sum = 0
    for num in groupArr {
        sum += worldMap[num.0][num.1]
    }
    let avr = sum / groupArr.count
    for num in groupArr {
        worldMap[num.0][num.1] = avr
    }
}

func initalIsGroupMap() {
    isGroupMap = Array(repeating: Array(repeating: false, count: n), count: n)
}

while true {
    var groupArr: [[(Int, Int)]] = []
    for i in 0..<n {
        for j in 0..<n {
            if isGroupMap[i][j] {
                continue
            }
            let group = makeGroup(i: i, j: j, group: [])
            if group.count > 1 {
                groupArr.append(group)
            }
        }
    }
    groupArr.forEach { arr in
        calculateWorldMap(groupArr: arr)
    }
    if groupArr.count == 0 {
        break
    }
    initalIsGroupMap()
    answer += 1
}

print(answer)