LeetCode in Kotlin

3501. Maximize Active Section with Trade II

Hard

You are given a binary string s of length n, where:

You can perform at most one trade to maximize the number of active sections in s. In a trade, you:

Additionally, you are given a 2D array queries, where queries[i] = [li, ri] represents a substring s[li...ri].

For each query, determine the maximum possible number of active sections in s after making the optimal trade on the substring s[li...ri].

Return an array answer, where answer[i] is the result for queries[i].

Note

Example 1:

Input: s = “01”, queries = [[0,1]]

Output: [1]

Explanation:

Because there is no block of '1's surrounded by '0's, no valid trade is possible. The maximum number of active sections is 1.

Example 2:

Input: s = “0100”, queries = [[0,3],[0,2],[1,3],[2,3]]

Output: [4,3,1,1]

Explanation:

Example 3:

Input: s = “1000100”, queries = [[1,5],[0,6],[0,4]]

Output: [6,7,2]

Explanation:

Example 4:

Input: s = “01010”, queries = [[0,3],[1,4],[1,3]]

Output: [4,4,2]

Explanation:

Constraints:

Solution

import kotlin.math.max

class Solution {
    fun maxActiveSectionsAfterTrade(s: String, queries: Array<IntArray>): List<Int> {
        val n = s.length
        var activeCount = 0
        for (ch in s.toCharArray()) {
            if (ch == '1') {
                activeCount++
            }
        }
        val segments: MutableList<IntArray?> = ArrayList<IntArray?>()
        var start = 0
        for (i in 0..<n) {
            if (i == n - 1 || s[i] != s[i + 1]) {
                segments.add(intArrayOf(start, i - start + 1))
                start = i + 1
            }
        }
        val segmentCount = segments.size
        val maxPower = 20
        val rmq = Array<IntArray>(maxPower) { IntArray(segmentCount) }
        for (i in 0..<maxPower) {
            rmq[i].fill(NEG_INF)
        }
        for (i in 0..<segmentCount) {
            if (s[segments[i]!![0]] == '0' && i + 2 < segmentCount) {
                rmq[0][i] = segments[i]!![1] + segments[i + 2]!![1]
            }
        }
        var power = 1
        var rangeLen = 2
        while (power < maxPower) {
            var i = 0
            var j = rangeLen - 1
            while (j < segmentCount) {
                rmq[power][i] =
                    max(rmq[power - 1][i], rmq[power - 1][i + rangeLen / 2])
                i++
                j++
            }
            power++
            rangeLen *= 2
        }
        val result: MutableList<Int> = ArrayList<Int>()
        for (query in queries) {
            val left = query[0]
            val right = query[1]
            val leftIndex = binarySearch(segments, left) - 1
            val rightIndex = binarySearch(segments, right) - 1
            if (rightIndex - leftIndex + 1 <= 2) {
                result.add(activeCount)
                continue
            }
            var bestIncrease = max(getMaxInRange(rmq, leftIndex + 1, rightIndex - 3), 0)
            bestIncrease = max(
                bestIncrease,
                calculateNewSections(
                    s, segments, left, right, leftIndex, rightIndex, leftIndex,
                ),
            )
            bestIncrease = max(
                bestIncrease,
                calculateNewSections(
                    s,
                    segments,
                    left,
                    right,
                    leftIndex,
                    rightIndex,
                    rightIndex - 2,
                ),
            )
            result.add(activeCount + bestIncrease)
        }
        return result
    }

    private fun binarySearch(segments: MutableList<IntArray?>, key: Int): Int {
        var lo = 0
        var hi = segments.size
        while (lo < hi) {
            val mid = lo + (hi - lo) / 2
            if (segments[mid]!![0] > key) {
                hi = mid
            } else {
                lo = mid + 1
            }
        }
        return lo
    }

    private fun getMaxInRange(rmq: Array<IntArray>, left: Int, right: Int): Int {
        if (left > right) {
            return NEG_INF
        }
        val power = 31 - Integer.numberOfLeadingZeros(right - left + 1)
        return max(rmq[power][left], rmq[power][right - (1 shl power) + 1])
    }

    private fun getSegmentSize(
        segments: MutableList<IntArray?>,
        left: Int,
        right: Int,
        leftIndex: Int,
        rightIndex: Int,
        i: Int,
    ): Int {
        if (i == leftIndex) {
            return segments[leftIndex]!![1] - (left - segments[leftIndex]!![0])
        }
        if (i == rightIndex) {
            return right - segments[rightIndex]!![0] + 1
        }
        return segments[i]!![1]
    }

    private fun calculateNewSections(
        s: String,
        segments: MutableList<IntArray?>,
        left: Int,
        right: Int,
        leftIndex: Int,
        rightIndex: Int,
        i: Int,
    ): Int {
        if (i < 0 || i + 2 >= segments.size || s[segments[i]!![0]] == '1') {
            return NEG_INF
        }
        val size1 = getSegmentSize(segments, left, right, leftIndex, rightIndex, i)
        val size2 = getSegmentSize(segments, left, right, leftIndex, rightIndex, i + 2)
        return size1 + size2
    }

    companion object {
        private const val INF = 1e9
        private const val NEG_INF: Int = -INF.toInt()
    }
}
This site is open source. Improve this page.