[코딩테스트] 🚶 인구이동 - 백준 16234

문제링크

난이도

골드 4

유형

구현, BFS

소요시간

1시간

회고

정석적인 문제였다. 최단시간 풀이를 한 사람은 큐를 직접 구현했던데, 단순 Array로 시간 초과가 나면 큐를 구현해서 풀어야겠다.

코드

func current() {
    let NLR = readLine()!.split(separator: " ").map { Int($0)! }
    let (N, L, R) = (NLR[0], NLR[1], NLR[2])

    var A = Array(repeating: Array(repeating: 0, count: N), count: N)

    (0..<N).forEach { r in
        let row = readLine()!.split(separator: " ").map { Int($0)! }
        A[r] = row
    }

    let dr = [0, 0, -1, 1]
    let dc = [-1, 1, 0, 0]

    var answer = 0

    // 전체 순회 너비우선 탐색

    // 무한 루프
    while true {
        var visited = Array(repeating: Array(repeating: false, count: N), count: N)
        var nextA = A
        var moved = false

        // 완전 탐색
        (0..<N).forEach { r in
            (0..<N).forEach { c in
                if visited[r][c] { return }

                var q: [(r: Int, c: Int)] = [(r, c)]
                var cells: [(r: Int, c: Int)] = [(r, c)]
                var head = 0
                var acc = A[r][c]

                visited[r][c] = true

                // BFS
                while head < q.count {
                    let (cr, cc) = q[head]
                    head += 1
                    
                    (0..<4).forEach { k in
                        let nr = cr + dr[k]
                        let nc = cc + dc[k]
                        
                        if nr < 0 || nr >= N || nc < 0 || nc >= N { return }
                        if visited[nr][nc] { return }

                        // 문제의 조건
                        if (L...R) ~= abs(A[cr][cc] - A[nr][nc]) {
                            visited[nr][nc] = true
                            q.append((nr, nc))
                            cells.append((nr, nc))
                            acc += A[nr][nc]
                        }
                    }
                }

                // 인구 평균
                if cells.count > 1 {
                    let avg = acc / cells.count
                    cells.forEach { cr, cc in
                        nextA[cr][cc] = avg
                    }
                    moved = true
                }
            }
        }

        if !moved {
            print(answer)
            return
        }

        A = nextA
        answer += 1
    }
}