[白俊17143-Kotlin]釣り王

import java.io.BufferedReader
import java.io.BufferedWriter

private lateinit var bufferedReader: BufferedReader
private lateinit var bufferedWriter: BufferedWriter

// 1: (-1, 0)
// 2: (1, 0)
// 3: (0, 1)
// 4: (0, -1)
val dx = listOf(0, 0, 0, 1, -1)
val dy = listOf(0, -1, 1, 0, 0)

sealed class Element
data class Shark(var r: Int, var c: Int, val s: Int, var d: Int, val z: Int) : Element()
class Empty : Element()

fun main() {
    bufferedReader = System.`in`.bufferedReader()
    bufferedWriter = System.out.bufferedWriter()

    // 1. get (r, c, m)
    val (r, c, m) = bufferedReader
        .readLine()
        .split(" ")
        .map { it.toInt() }

    // 2. get shark's info
    val board: Array<Array<Element>> = Array(r + 1) { Array(c + 1) { Empty() } }
    val sharks = mutableListOf<Shark>()
    for (i in 0 until m) {
        val (r, c, s, d, z) = bufferedReader
            .readLine()
            .split(" ")
            .map { it.toInt() }
        val shark = Shark(r, c, s, d, z)
        board[r][c] = shark
        sharks.add(shark)
    }

    // 3. catch shark
    var total = 0
    for (i in 1..c) {
        total += catchShark(board, sharks, r, i)
        moveShark(board, sharks, r, c)
    }

    bufferedWriter.write("$total")

    bufferedReader.close()
    bufferedWriter.close()
}

fun catchShark(board: Array<Array<Element>>, sharks: MutableList<Shark>, r: Int, fisher: Int): Int {
    for (i in 1..r) {
        if (board[i][fisher] is Shark) {
            val catchedShark = board[i][fisher] as Shark
            board[i][fisher] = Empty()
            sharks.remove(catchedShark)
            return catchedShark.z
        }
    }
    return 0
}

fun moveShark(board: Array<Array<Element>>, sharks: MutableList<Shark>, r: Int, c: Int) {
    sharks.forEach { shark ->
        var remain = shark.s
        while (remain != 0) {
            val nx = remain * dx[shark.d] + shark.c
            val ny = remain * dy[shark.d] + shark.r
            if (nx > c) {
                remain -= (c - shark.c)
                shark.d = 4
                shark.c = c
            } else if (nx < 1) {
                remain -= (shark.c - 1)
                shark.d = 3
                shark.c = 1
            } else if (ny > r) {
                remain -= (r - shark.r)
                shark.d = 1
                shark.r = r
            } else if (ny < 1) {
                remain -= (shark.r - 1)
                shark.d = 2
                shark.r = 1
            } else {
                remain = 0
                shark.c = nx
                shark.r = ny
            }
        }
    }
    initBoard(board, r, c)
    checkEatenSharks(board, sharks)
}

fun initBoard(board: Array<Array<Element>>, r: Int, c: Int) {
    for (i in 1..r) {
        for (j in 1..c) {
            board[i][j] = Empty()
        }
    }
}

fun checkEatenSharks(board: Array<Array<Element>>, sharks: MutableList<Shark>) {
    val eatenSharks = mutableListOf<Shark>()
    sharks.forEach { shark ->
        if (board[shark.r][shark.c] is Shark) {
            val targetShark = board[shark.r][shark.c] as Shark
            if (targetShark.z < shark.z) {
                board[shark.r][shark.c] = shark
                eatenSharks.add(targetShark)
            } else eatenSharks.add(shark)
        } else board[shark.r][shark.c] = shark
    }
    sharks.removeAll(eatenSharks)
}