Условие

Вася очень любит везде искать своё счастливое число K. Каждый день он ходит в школу по улице, вдоль которой припарковано N машин. Он заинтересовался вопросом, сколько существует отрезков из подряд идущих машин таких, что сумма их номеров равна K. Помогите Васе узнать ответ на его вопрос.

Примеры

Ввод: 5 17 17 7 10 7 10 Вывод: 4

Ввод: 5 10 1 2 3 4 1 Вывод: 2

Решение

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))

    // Читаем количество элементов в массиве (n) и целевую сумму (k)
    val (n, k) = reader.readLine().split(" ").map { it.toInt() }
    // Читаем массив чисел
    val arr = reader.readLine().split(" ").map { it.toInt() }

    var currentSum = 0L // Текущая сумма префикса
    var count = 0 // Количество подмассивов с суммой k
    val prefixSumCount = mutableMapOf<Long, Int>() // Хранение количества встреченных сумм префиксов
    prefixSumCount[0L] = 1 // Начальная сумма 0 встречается 1 раз

    for (num in arr) {
        currentSum += num.toLong() // Обновляем текущую сумму префикса
        val target = currentSum - k.toLong() // Вычисляем разницу с целевой суммой
        count += prefixSumCount.getOrDefault(target, 0) // Добавляем количество найденных подмассивов
        prefixSumCount[currentSum] = prefixSumCount.getOrDefault(currentSum, 0) + 1 // Обновляем счетчик суммы префикса
    }

    writer.write(count.toString()) // Выводим результат
    writer.newLine()

    reader.close()
    writer.close()
}