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 7546988e..3544457b 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,208 +1,117 @@ package g3201_3300.s3245_alternating_groups_iii -// #Hard #Array #Binary_Indexed_Tree #2025_02_12_Time_188_ms_(100.00%)_Space_95.34_MB_(100.00%) +// #Hard #Array #Binary_Indexed_Tree #2025_03_15_Time_70_ms_(100.00%)_Space_108.25_MB_(_%) -import java.util.TreeMap -import kotlin.math.max +import java.util.BitSet class Solution { - // 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 - } - } - - // 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 + fun numberOfAlternatingGroups(colors: IntArray, queries: Array): MutableList { + val n = colors.size + val set = BitSet() + val bit = BIT(n) + for (i in 0.. = ArrayList() + for (q in queries) { + if (q[0] == 1) { + if (set.isEmpty) { + ans.add(n) + } else { + val size = q[1] + val res = bit.query(size) + ans.add(res[1] - res[0] * (size - 1)) + } + } else { + val i = q[1] + var color = colors[i] + if (q[2] == color) { + continue + } + val pre = getIndex(i - 1, n) + if (colors[pre] == color) { + remove(set, bit, n, pre) + } + val next = getIndex(i + 1, n) + if (colors[next] == color) { + remove(set, bit, n, i) + } + colors[i] = colors[i] xor 1 + color = colors[i] + if (colors[pre] == color) { + add(set, bit, n, pre) + } + if (colors[next] == color) { + add(set, bit, n, i) + } } } + 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) - } - - // 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) - } - - // Returns the length of a group from index x to y. - private fun len(x: Int, y: Int): Int { - return y - x + 1 - } - - // --- 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) - } - - // Adds a group (block) by updating BIT with a positive value. - private fun addGroup(start: Int, end: Int) { - edt(len(start, end), 1) - } - - // 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 + private fun add(set: BitSet, bit: BIT, n: Int, i: Int) { + if (set.isEmpty) { + bit.update(n, 1) + } else { + update(set, bit, n, i, 1) } + set.set(i) } - // 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() + private fun remove(set: BitSet, bit: BIT, n: Int, i: Int) { + set.clear(i) + if (set.isEmpty) { + bit.update(n, -1) + } else { + update(set, bit, n, i, -1) } - return ans } - // 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 + private fun update(set: BitSet, bit: BIT, n: Int, i: Int, v: Int) { + var pre = set.previousSetBit(i) + if (pre == -1) { + pre = set.previousSetBit(n) } - // 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) - } - } + var next = set.nextSetBit(i) + if (next == -1) { + next = set.nextSetBit(0) } - // 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) - } - } - } - - // Merge the affected groups into one new group. - groups.put(ml, mr) - addGroup(ml, mr) + bit.update(getIndex(next - pre + n - 1, n) + 1, -v) + bit.update(getIndex(i - pre, n), v) + bit.update(getIndex(next - i, n), v) } - // 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) - } + private fun getIndex(index: Int, mod: Int): Int { + val result = if (index >= mod) index - mod else index + return if (index < 0) index + mod else result } - // 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) + private class BIT(n: Int) { + var n: Int = n + 1 + var tree1: IntArray = IntArray(n + 1) + var tree2: IntArray = IntArray(n + 1) + + fun update(size: Int, v: Int) { + var i = size + while (i > 0) { + tree1[i] += v + tree2[i] += v * size + i -= i and -i } } - // 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()) + fun query(size: Int): IntArray { + var count = 0 + var sum = 0 + var i = size + while (i < n) { + count += tree1[i] + sum += tree2[i] + i += i and -i + } + return intArrayOf(count, sum) + } } }