diff --git a/src/main/kotlin/g3201_3300/s3245_alternating_groups_iii/Solution.kt b/src/main/kotlin/g3201_3300/s3245_alternating_groups_iii/Solution.kt index 6dbff5331..cf9ea0c04 100644 --- a/src/main/kotlin/g3201_3300/s3245_alternating_groups_iii/Solution.kt +++ b/src/main/kotlin/g3201_3300/s3245_alternating_groups_iii/Solution.kt @@ -1,260 +1,208 @@ package g3201_3300.s3245_alternating_groups_iii -// #Hard #Array #Binary_Indexed_Tree #2024_08_07_Time_1072_ms_(100.00%)_Space_97.6_MB_(100.00%) +// #Hard #Array #Binary_Indexed_Tree #2025_02_12_Time_188_ms_(100.00%)_Space_95.34_MB_(100.00%) + +import java.util.TreeMap +import kotlin.math.max -@Suppress("NAME_SHADOWING") class Solution { - private fun go(ind: Int, lst: LST, fs: IntArray, n: Int, ff: LST, c: IntArray) { - if (ind > 0) { - val pre = lst.prev(ind - 1) - var nex = lst.next(pre + 1) - if (nex == -1) { - nex = 2 * n - } - if (pre != -1 && pre < n && --fs[nex - pre] == 0) { - ff.unsetPos(nex - pre) - } - } - if (lst.get(ind)) { - val pre = ind - var nex = lst.next(ind + 1) - if (nex == -1) { - nex = 2 * n - } - if (pre != -1 && pre < n && --fs[nex - pre] == 0) { - ff.unsetPos(nex - pre) - } - } - if (lst.get(ind + 1)) { - val pre = ind + 1 - var nex = lst.next(ind + 2) - if (nex == -1) { - nex = 2 * n - } - if (pre != -1 && pre < n && --fs[nex - pre] == 0) { - ff.unsetPos(nex - pre) - } - } - lst.unsetPos(ind) - lst.unsetPos(ind + 1) - c[ind] = c[ind] xor 1 - if (ind > 0 && c[ind] != c[ind - 1]) { - lst.setPos(ind) - } - if (ind + 1 < c.size && c[ind + 1] != c[ind]) { - lst.setPos(ind + 1) - } - if (ind > 0) { - val pre = lst.prev(ind - 1) - var nex = lst.next(pre + 1) - if (nex == -1) { - nex = 2 * n - } - if (pre != -1 && pre < n && ++fs[nex - pre] == 1) { - ff.setPos(nex - pre) + // Binary Indexed Tree (BIT) class. + private class BIT { + var bs: IntArray = IntArray(SZ) + + // Update BIT: add value y to index x. + fun update(x: Int, y: Int) { + var x = x + x = OFFSET - x + while (x < SZ) { + bs[x] += y + x += x and -x } } - if (lst.get(ind)) { - val pre = ind - var nex = lst.next(ind + 1) - if (nex == -1) { - nex = 2 * n - } - if (pre < n && ++fs[nex - pre] == 1) { - ff.setPos(nex - pre) + + // Query BIT: get the prefix sum up to index x. + fun query(x: Int): Int { + var x = x + x = OFFSET - x + var ans = 0 + while (x > 0) { + ans += bs[x] + x -= x and -x } + return ans } - if (lst.get(ind + 1)) { - val pre = ind + 1 - var nex = lst.next(ind + 2) - if (nex == -1) { - nex = 2 * n - } - if (pre < n && ++fs[nex - pre] == 1) { - ff.setPos(nex - pre) + + // Clear BIT values starting from index x. + fun clear(x: Int) { + var x = x + x = OFFSET - x + while (x < SZ) { + bs[x] = 0 + x += x and -x } } } - fun numberOfAlternatingGroups(colors: IntArray, queries: Array): List { - val n = colors.size - val c = IntArray(2 * n) - for (i in 0 until 2 * n) { - c[i] = colors[i % n] xor (if (i % 2 == 0) 0 else 1) - } - val lst = LST(2 * n + 3) - for (i in 1 until 2 * n) { - if (c[i] != c[i - 1]) { - lst.setPos(i) - } - } - val fs = IntArray(2 * n + 1) - val ff = LST(2 * n + 1) - for (i in 0 until n) { - if (lst.get(i)) { - var ne = lst.next(i + 1) - if (ne == -1) { - ne = 2 * n - } - fs[ne - i]++ - ff.setPos(ne - i) - } - } - val ans: MutableList = ArrayList() - for (q in queries) { - if (q[0] == 1) { - if (lst.next(0) == -1) { - ans.add(n) - } else { - var lans = 0 - var i = ff.next(q[1]) - while (i != -1) { - lans += (i - q[1] + 1) * fs[i] - i = ff.next(i + 1) - } - if (c[2 * n - 1] != c[0]) { - val f = lst.next(0) - if (f >= q[1]) { - lans += (f - q[1] + 1) - } - } - ans.add(lans) - } - } else { - val ind = q[1] - val `val` = q[2] - if (colors[ind] == `val`) { - continue - } - colors[ind] = colors[ind] xor 1 - go(ind, lst, fs, n, ff, c) - go(ind + n, lst, fs, n, ff, c) - } - } - return ans + // --- BIT wrapper methods --- + // Updates both BITs for a given group length. + private fun edt(x: Int, y: Int) { + // Second BIT is updated with x * y. + BITS[1].update(x, x * y) + // First BIT is updated with y. + BITS[0].update(x, y) } - private class LST(private val n: Int) { - private val set: Array + // Combines BIT queries to get the result for a given x. + private fun qry(x: Int): Int { + return BITS[1].query(x) + (1 - x) * BITS[0].query(x) + } - init { - var d = 1 - d = getD(n, d) - set = arrayOfNulls(d) - var i = 0 - var m = n ushr 6 - while (i < d) { - set[i] = LongArray(m + 1) - i++ - m = m ushr 6 - } - } + // Returns the length of a group from index x to y. + private fun len(x: Int, y: Int): Int { + return y - x + 1 + } - private fun getD(n: Int, d: Int): Int { - var d = d - var m = n - while (m > 1) { - m = m ushr 6 - d++ - } - return d - } + // --- Group operations --- + // Removes a group (block) by updating BIT with a negative value. + private fun removeGroup(start: Int, end: Int) { + edt(len(start, end), -1) + } - fun setPos(pos: Int): LST { - var pos = pos - if (pos >= 0 && pos < n) { - var i = 0 - while (i < set.size) { - set[i]!![pos ushr 6] = set[i]!![pos ushr 6] or (1L shl pos) - i++ - pos = pos ushr 6 - } - } - return this - } + // Adds a group (block) by updating BIT with a positive value. + private fun addGroup(start: Int, end: Int) { + edt(len(start, end), 1) + } - fun unsetPos(pos: Int): LST { - var pos = pos - if (pos >= 0 && pos < n) { - var i = 0 - while (i < set.size && (i == 0 || set[i - 1]!![pos] == 0L) - ) { - set[i]!![pos ushr 6] = set[i]!![pos ushr 6] and (1L shl pos).inv() - i++ - pos = pos ushr 6 - } - } - return this + // Initializes the alternating groups using the colors array. + private fun initializeGroups(colors: IntArray, groups: TreeMap) { + val n = colors.size + var i = 0 + while (i < n) { + var r = i + // Determine the end of the current alternating group. + while (r < n && (colors[r] + colors[i] + r + i) % 2 == 0) { + ++r + } + // The group spans from index i to r-1. + groups.put(i, r - 1) + // Update BITs with the length of this group. + edt(r - i, 1) + // Skip to the end of the current group. + i = r - 1 + ++i } + } - fun get(pos: Int): Boolean { - return pos >= 0 && pos < n && set[0]!![pos ushr 6] shl pos.inv() < 0 + // Processes a type 1 query: returns the number of alternating groups + // of at least the given size. + private fun processQueryType1(colors: IntArray, groups: TreeMap, groupSize: Int): Int { + var ans = qry(groupSize) + val firstGroup = groups.firstEntry() + val lastGroup = groups.lastEntry() + // If there is more than one group and the first and last colors differ, + // adjust the answer by "merging" the groups at the boundaries. + if (firstGroup !== lastGroup && colors[0] != colors[colors.size - 1]) { + val leftLen = len(firstGroup.key!!, firstGroup.value!!) + val rightLen = len(lastGroup.key!!, lastGroup.value!!) + ans = (ans - max((leftLen - groupSize + 1).toDouble(), 0.0)).toInt() + ans = (ans - max((rightLen - groupSize + 1).toDouble(), 0.0)).toInt() + ans = (ans + max((leftLen + rightLen - groupSize + 1).toDouble(), 0.0)).toInt() } + return ans + } - fun prev(pos: Int): Int { - var pos = pos - var i = 0 - while (i < set.size && pos >= 0) { - val pre = prev(set[i]!![pos ushr 6], pos and 63) - if (pre != -1) { - pos = pos ushr 6 shl 6 or pre - while (i > 0) { - pos = pos shl 6 or 63 - java.lang.Long.numberOfLeadingZeros(set[--i]!![pos]) - } - return pos + // Processes a type 2 query: updates the color at index x and adjusts groups. + private fun processQueryType2( + colors: IntArray, + groups: TreeMap, + x: Int, + newColor: Int, + ) { + if (colors[x] == newColor) { + return + } + // Update the color at index x. + colors[x] = newColor + // Find the group (block) that contains index x. + var it = groups.floorEntry(x) + val l: Int = it!!.key!! + val r: Int = it.value!! + // Remove the old group from BIT and map. + removeGroup(l, r) + groups.remove(l) + var ml = x + var mr = x + // Process the left side of index x. + if (l != x) { + groups.put(l, x - 1) + addGroup(l, x - 1) + } else { + if (x > 0 && colors[x] != colors[x - 1]) { + it = groups.floorEntry(x - 1) + if (it != null) { + ml = it.key!! + removeGroup(it.key!!, it.value!!) + groups.remove(it.key) } - i++ - pos = pos ushr 6 - pos-- - } - return -1 - } - - private fun prev(set: Long, n: Int): Int { - val h = set shl n.inv() - if (h == 0L) { - return -1 } - return -java.lang.Long.numberOfLeadingZeros(h) + n } - - fun next(pos: Int): Int { - var pos = pos - var i = 0 - while (i < set.size && pos ushr 6 < set[i]!!.size) { - val nex = next(set[i]!![pos ushr 6], pos and 63) - if (nex != -1) { - pos = pos ushr 6 shl 6 or nex - while (i > 0) { - pos = pos shl 6 or java.lang.Long.numberOfTrailingZeros(set[--i]!![pos]) - } - return pos + // Process the right side of index x. + if (r != x) { + groups.put(x + 1, r) + addGroup(x + 1, r) + } else { + if (x + 1 < colors.size && colors[x + 1] != colors[x]) { + it = groups.ceilingEntry(x + 1) + if (it != null) { + mr = it.value!! + removeGroup(it.key!!, it.value!!) + groups.remove(it.key) } - i++ - pos = pos ushr 6 - pos++ } - return -1 } - override fun toString(): String { - val list: MutableList = ArrayList() - var pos = next(0) - while (pos != -1) { - list.add(pos) - pos = next(pos + 1) - } - return list.toString() + // Merge the affected groups into one new group. + groups.put(ml, mr) + addGroup(ml, mr) + } + + // Clears both BITs. This is done after processing all queries. + private fun clearAllBITs(n: Int) { + for (i in 0..n + 2) { + BITS[0].clear(i) + BITS[1].clear(i) } + } - companion object { - private fun next(set: Long, n: Int): Int { - val h = set ushr n - if (h == 0L) { - return -1 - } - return java.lang.Long.numberOfTrailingZeros(h) + n + // Main function to handle queries on alternating groups. + fun numberOfAlternatingGroups(colors: IntArray, queries: Array): MutableList { + val groups = TreeMap() + val n = colors.size + val results: MutableList = ArrayList() + // Initialize alternating groups. + initializeGroups(colors, groups) + // Process each query. + for (query in queries) { + if (query[0] == 1) { + // Type 1 query: count alternating groups of at least a given size. + val groupSize = query[1] + val ans = processQueryType1(colors, groups, groupSize) + results.add(ans) + } else { + // Type 2 query: update the color at a given index. + val index = query[1] + val newColor = query[2] + processQueryType2(colors, groups, index, newColor) } } + // Clear BITs after processing all queries. + clearAllBITs(n) + return results + } + + companion object { + private const val SZ = 63333 + private const val OFFSET: Int = SZ - 10 + private val BITS = arrayOf(BIT(), BIT()) } } diff --git a/src/test/kotlin/g3201_3300/s3245_alternating_groups_iii/SolutionTest.kt b/src/test/kotlin/g3201_3300/s3245_alternating_groups_iii/SolutionTest.kt index 2518012fb..6123835d5 100644 --- a/src/test/kotlin/g3201_3300/s3245_alternating_groups_iii/SolutionTest.kt +++ b/src/test/kotlin/g3201_3300/s3245_alternating_groups_iii/SolutionTest.kt @@ -28,4 +28,16 @@ internal class SolutionTest { equalTo(listOf(2, 0)), ) } + + @Test + fun numberOfAlternatingGroups3() { + assertThat( + Solution() + .numberOfAlternatingGroups( + intArrayOf(0, 0, 0, 1), + arrayOf(intArrayOf(2, 1, 1), intArrayOf(1, 3), intArrayOf(2, 1, 1), intArrayOf(2, 0, 1)), + ), + equalTo(listOf(4)), + ) + } }