diff --git a/kt/student-attendance-record-ii.kt b/kt/student-attendance-record-ii.kt new file mode 100644 index 0000000..27ed5d1 --- /dev/null +++ b/kt/student-attendance-record-ii.kt @@ -0,0 +1,43 @@ +val MOD = 1000000007 + +class Solution { + private fun createState(): MutableList> = (0..1).map { mutableListOf(0, 0, 0) }.toMutableList() + + fun checkRecord(n: Int): Int { + val initialState = createState() + initialState[0][0] = 1 + + val finalState = + (1..n).fold(initialState) { state, _ -> + val nextState = createState() + + (0..1).flatMap { absences -> (0..2).map { lates -> absences to lates } }.forEach { (absences, lates) -> + nextState[absences][0] = (nextState[absences][0] + state[absences][lates]) % MOD + + if (absences < 1) { + nextState[absences + 1][0] = (nextState[absences + 1][0] + state[absences][lates]) % MOD + } + + if (lates < 2) { + nextState[absences][lates + 1] = (nextState[absences][lates + 1] + state[absences][lates]) % MOD + } + } + + return@fold nextState + } + + return finalState.fold(0) { acc, it -> + it.fold(acc) { acc, it -> + (acc + it) % MOD + } + } + } +} + +fun main() { + val s = Solution() + + check(s.checkRecord(2) == 8) + check(s.checkRecord(1) == 3) + check(s.checkRecord(10101) == 183236316) +}