Решение
import java.io.BufferedReader
import java.io.BufferedWriter
import java.io.InputStreamReader
import java.io.OutputStreamWriter
import kotlin.math.floor
import kotlin.math.sqrt
fun main(args: Array<String>) {
val reader = BufferedReader(InputStreamReader(System.`in`))
val writer = BufferedWriter(OutputStreamWriter(System.out))
val n = reader.readLine().toInt()
val tokens = reader.readLine().split(" ")
val a = IntArray(n)
for (i in 0 until n) {
a[i] = tokens[i].toInt()
}
// Нахождение максимального элемента, чтобы вычислить K_max
var maxA = 0
for (value in a) {
if (value > maxA) maxA = value
}
// Если maxA == 0, ни один подотрезок не сможет удовлетворить условию для k>=1
val K_max = if (maxA > 0) floor(sqrt(maxA.toDouble())).toInt() else 0
var result = 0L
// Для каждого k от 1 до K_max считаем число подотрезков,
// где количество элементов с a[i] >= k^2 не меньше k.
for (k in 1..K_max) {
val thresh = k * k
var sumWindow = 0
var r = 0
// Применяем технику скользящего окна для подсчёта подотрезков,
// удовлетворяющих условию: сумма по массиву b = [a[i]>=thresh] должна быть >= k.
for (l in 0 until n) {
while (r < n && sumWindow < k) {
if (a[r] >= thresh) {
sumWindow++
}
r++
}
if (sumWindow >= k) {
// Если условие выполнено, то для фиксированного l все подотрезки [l, r'] с r' >= r удовлетворяют условию.
result += (n - r + 1)
} else {
// Если для текущего l не удалось набрать k, то для всех больших l также не наберем,
// поэтому можно прерывать цикл по l.
break
}
if (a[l] >= thresh) {
sumWindow--
}
}
}
writer.write(result.toString())
writer.flush()
writer.close()
reader.close()
}