https://leetcode.com/problems/distribute-candies-to-people Easy

Решение

class Solution {
    fun distributeCandies(candies: Int, num_people: Int): IntArray {
        val c = candies.toLong()
        val n = num_people
        var k = (((Math.sqrt(1.0 + 8.0 * c) - 1.0) / 2.0)).toLong()
        var sumK = k * (k + 1) / 2
        while (sumK > c) {
            k--
            sumK = k * (k + 1) / 2
        }
        while ((k + 1) * (k + 2) / 2 <= c) {
            k++
            sumK = k * (k + 1) / 2
        }
        val full = (k / n).toLong()
        val rem = (k % n).toInt()
        val res = IntArray(n)
        val nL = n.toLong()
        var i = 0
        while (i < n) {
            val cnt = full + if (i < rem) 1L else 0L
            if (cnt > 0L) {
                val first = (i + 1).toLong()
                val last = first + (cnt - 1L) * nL
                val sum = (first + last) * cnt / 2L
                res[i] = sum.toInt()
            }
            i++
        }
        val leftover = (c - sumK).toInt()
        if (leftover > 0 && rem < n) res[rem] += leftover
        return res
    }
}