LeetCode in Kotlin

3003. Maximize the Number of Partitions After Operations

Hard

You are given a 0-indexed string s and an integer k.

You are to perform the following partitioning operations until s is empty:

Before the operations, you are allowed to change at most one index in s to another lowercase English letter.

Return an integer denoting the maximum number of resulting partitions after the operations by optimally choosing at most one index to change.

Example 1:

Input: s = “accca”, k = 2

Output: 3

Explanation: In this example, to maximize the number of resulting partitions, s[2] can be changed to ‘b’. s becomes “acbca”. The operations can now be performed as follows until s becomes empty:

Hence, the answer is 3. It can be shown that it is not possible to obtain more than 3 partitions.

Example 2:

Input: s = “aabaab”, k = 3

Output: 1

Explanation: In this example, to maximize the number of resulting partitions we can leave s as it is. The operations can now be performed as follows until s becomes empty:

Hence, the answer is 1. It can be shown that it is not possible to obtain more than 1 partition.

Example 3:

Input: s = “xxyz”, k = 1

Output: 4

Explanation: In this example, to maximize the number of resulting partitions, s[1] can be changed to ‘a’. s becomes “xayz”. The operations can now be performed as follows until s becomes empty:

Hence, the answer is 4. It can be shown that it is not possible to obtain more than 4 partitions.

Constraints:

Solution

import kotlin.math.max

class Solution {
    fun maxPartitionsAfterOperations(s: String, k: Int): Int {
        if (k == ALPHABET_SIZE) {
            return 1
        }
        val n = s.length
        val ansr = IntArray(n)
        val usedr = IntArray(n)
        var used = 0
        var cntUsed = 0
        var ans = 1
        for (i in n - 1 downTo 0) {
            val ch = s[i].code - 'a'.code
            if ((used and (1 shl ch)) == 0) {
                if (cntUsed == k) {
                    cntUsed = 0
                    used = 0
                    ans++
                }
                used = used or (1 shl ch)
                cntUsed++
            }
            ansr[i] = ans
            usedr[i] = used
        }
        var ansl = 0
        ans = ansr[0]
        var l = 0
        while (l < n) {
            used = 0
            cntUsed = 0
            var usedBeforeLast = 0
            var usedTwiceBeforeLast = 0
            var last = -1
            var r = l
            while (r < n) {
                val ch = s[r].code - 'a'.code
                if ((used and (1 shl ch)) == 0) {
                    if (cntUsed == k) {
                        break
                    }
                    usedBeforeLast = used
                    last = r
                    used = used or (1 shl ch)
                    cntUsed++
                } else if (cntUsed < k) {
                    usedTwiceBeforeLast = usedTwiceBeforeLast or (1 shl ch)
                }
                r++
            }
            if (cntUsed == k) {
                if (last - l > Integer.bitCount(usedBeforeLast)) {
                    ans = max(ans, (ansl + 1 + ansr[last]))
                }
                if (last + 1 < r) {
                    if (last + 2 >= n) {
                        ans = max(ans, (ansl + 1 + 1))
                    } else {
                        if (Integer.bitCount(usedr[last + 2]) == k) {
                            val canUse = ((1 shl ALPHABET_SIZE) - 1) and used.inv() and usedr[last + 2].inv()
                            ans = if (canUse > 0) {
                                max(ans, (ansl + 1 + 1 + ansr[last + 2]))
                            } else {
                                max(ans, (ansl + 1 + ansr[last + 2]))
                            }
                            val l1 = s[last + 1].code - 'a'.code
                            if ((usedTwiceBeforeLast and (1 shl l1)) == 0) {
                                ans = max(ans, (ansl + 1 + ansr[last + 1]))
                            }
                        } else {
                            ans = max(ans, (ansl + 1 + ansr[last + 2]))
                        }
                    }
                }
            }
            l = r
            ansl++
        }
        return ans
    }

    companion object {
        private const val ALPHABET_SIZE = 'z'.code - 'a'.code + 1
    }
}