educative.io

Educative

Kotlin using one heap, is it correct?

fun findMaximumCapital(
        capital: IntArray,
        profits: IntArray,
        numberOfProjects: Int,
        initialCapital: Int
    ): Int {
        val maxHeap = PriorityQueue<Pair<Int, Int>> { a, b -> b.second - a.second }
        var result = initialCapital
        // initial filling
        var i = 0
        i = fillNewProjects(i, capital, result, maxHeap, profits)
        for (j in 0 until numberOfProjects) { // N*logN + K*logN
            val poll = maxHeap.poll()
            result += poll.second
            i = fillNewProjects(i, capital, result, maxHeap, profits)
        }

        return result
    }

    private fun fillNewProjects(
        i: Int,
        capital: IntArray,
        initialCapital: Int,
        maxHeap: PriorityQueue<Pair<Int, Int>>,
        profits: IntArray
    ): Int {
        var i = i
        while (i < capital.size && capital[i] <= initialCapital) {
            maxHeap.offer(Pair(capital[i], profits[i]))
            i++
        }
        return i - 1
    }