백준 34060번 로봇 청소기 Kotlin

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

문제

---

생각하기

---

• 유니온 파인드 문제이다.

동작

결과

코드

---

import java.io.BufferedWriter
import java.io.File
import java.io.OutputStreamWriter

// input
private var br = System.`in`.bufferedReader()

// variables
private var N = 0
private lateinit var rank: IntArray
private lateinit var parent: IntArray
private lateinit var y: IntArray
private var components = 0

fun main() {
    val bw = BufferedWriter(OutputStreamWriter(System.out))

    input()
    bw.write(solve())
    bw.close()
} // End of main()

private fun solve(): String {
    val sb = StringBuilder()

    val prefixBreaks = IntArray(N)
    for (i in 1 until N) {
        prefixBreaks[i] = prefixBreaks[i - 1]
        if (y[i - 1] >= y[i]) {
            prefixBreaks[i]++
        }
    }

    val horizontalMap = HashMap<Int, HashMap<Int, Int>>()
    val verticalMap = HashMap<Int, HashMap<Int, Int>>()
    for (i in 0 until N) {
        val curY = y[i]
        val curBreakCount = prefixBreaks[i]

        if (horizontalMap.containsKey(curY)) {
            val breakToIdxMap = horizontalMap.get(curY)
            if (breakToIdxMap!!.containsKey(curBreakCount)) {
                union(i, breakToIdxMap[curBreakCount]!!)
            }

            if (breakToIdxMap.containsKey(curBreakCount - 1)) {
                union(i, breakToIdxMap[curBreakCount - 1]!!)
            }
        }

        if (verticalMap.containsKey(curBreakCount)) {
            val yToIdxMap = verticalMap[curBreakCount]
            if (yToIdxMap!!.containsKey(curY - 1)) {
                union(i, yToIdxMap[curY - 1]!!)
            }
            if (yToIdxMap.containsKey(curY + 1)) {
                union(i, yToIdxMap[curY + 1]!!)
            }
        }

        horizontalMap.getOrPut(curY) { HashMap() }[curBreakCount] = i
        verticalMap.getOrPut(curBreakCount) { HashMap() }[curY] = i
    }


    sb.append(components).append('\n').append(N)

    return sb.toString()
} // End of solve()

private fun find(root: Int): Int {
    if (root == parent[root]) return root

    parent[root] = find(parent[root])
    return parent[root]
} // End of find()

private fun union(a: Int, b: Int) {
    val rootA = find(a)
    val rootB = find(b)

    if (rootA != rootB) {
        if (rank[rootA] < rank[rootB]) {
            parent[rootA] = rootB
        } else if (rank[rootA] > rank[rootB]) {
            parent[rootB] = rootA
        } else {
            parent[rootA] = rootB
            rank[rootB]++
        }
        components--
    }
} // End of union()

private fun input() {
    N = br.readLine().toInt()

    y = IntArray(N) {
        br.readLine().toInt()
    }

    parent = IntArray(N) { it }
    rank = IntArray(N) { 0 }
    components = N
} // End of input()