https://coderun.yandex.ru/problem/robot/description Средняя

Решение

import java.io.BufferedReader
import java.io.BufferedWriter
import java.io.InputStreamReader
import java.io.OutputStreamWriter

fun main() {
    val reader = BufferedReader(InputStreamReader(System.`in`))
    val writer = BufferedWriter(OutputStreamWriter(System.out))
    
    val K = reader.readLine().trim().toInt()
    val process = reader.readLine().trim()
    val N = process.length

    if(K == 0) {
        val total = N.toLong() * (N.toLong() + 1) / 2
        writer.write(total.toString())
        writer.flush()
        reader.close()
        writer.close()
        return
    }
    
    val sizeB = N - K
    val B = BooleanArray(sizeB) { i -> process[i] == process[i + K] }
    
    val nxt = IntArray(sizeB)
    nxt[sizeB - 1] = if (!B[sizeB - 1]) sizeB - 1 else sizeB
    for(i in sizeB - 2 downTo 0) {
        nxt[i] = if (!B[i]) i else nxt[i + 1]
    }
    
    var ans = 0L
    for(L in 0 until sizeB) {
        val f = nxt[L] - L
        val lowR = L + K
        val highR = (L + f + K - 1).coerceAtMost(N - 1)
        if(lowR <= highR) {
            ans += (highR - lowR + 1)
        }
    }
    
    writer.write(ans.toString())
    writer.flush()
    reader.close()
    writer.close()
}