Skip to content

Commit 1161b7a

Browse files
javadevjavadev
committed
Improved task 3245
1 parent 6acc1d5 commit 1161b7a

File tree

1 file changed

+89
-180
lines changed
  • src/main/kotlin/g3201_3300/s3245_alternating_groups_iii

1 file changed

+89
-180
lines changed

‎src/main/kotlin/g3201_3300/s3245_alternating_groups_iii/Solution.kt

+89-180
Original file line numberDiff line numberDiff line change
@@ -1,208 +1,117 @@
11
package g3201_3300.s3245_alternating_groups_iii
22

3-
// #Hard #Array #Binary_Indexed_Tree #2025_02_12_Time_188_ms_(100.00%)_Space_95.34_MB_(100.00%)
3+
// #Hard #Array #Binary_Indexed_Tree #2025_03_15_Time_70_ms_(100.00%)_Space_108.25_MB_(_%)
44

5-
import java.util.TreeMap
6-
import kotlin.math.max
5+
import java.util.BitSet
76

87
class Solution {
9-
// Binary Indexed Tree (BIT) class.
10-
private class BIT {
11-
var bs: IntArray = IntArray(SZ)
12-
13-
// Update BIT: add value y to index x.
14-
fun update(x: Int, y: Int) {
15-
var x = x
16-
x = OFFSET - x
17-
while (x < SZ) {
18-
bs[x] += y
19-
x += x and -x
20-
}
21-
}
22-
23-
// Query BIT: get the prefix sum up to index x.
24-
fun query(x: Int): Int {
25-
var x = x
26-
x = OFFSET - x
27-
var ans = 0
28-
while (x > 0) {
29-
ans += bs[x]
30-
x -= x and -x
8+
fun numberOfAlternatingGroups(colors: IntArray, queries: Array<IntArray>): MutableList<Int?> {
9+
val n = colors.size
10+
val set = BitSet()
11+
val bit = BIT(n)
12+
for (i in 0..<n) {
13+
if (colors[i] == colors[getIndex(i + 1, n)]) {
14+
add(set, bit, n, i)
3115
}
32-
return ans
3316
}
34-
35-
// Clear BIT values starting from index x.
36-
fun clear(x: Int) {
37-
var x = x
38-
x = OFFSET - x
39-
while (x < SZ) {
40-
bs[x] = 0
41-
x += x and -x
17+
val ans: MutableList<Int?> = ArrayList<Int?>()
18+
for (q in queries) {
19+
if (q[0] == 1) {
20+
if (set.isEmpty) {
21+
ans.add(n)
22+
} else {
23+
val size = q[1]
24+
val res = bit.query(size)
25+
ans.add(res[1] - res[0] * (size - 1))
26+
}
27+
} else {
28+
val i = q[1]
29+
var color = colors[i]
30+
if (q[2] == color) {
31+
continue
32+
}
33+
val pre = getIndex(i - 1, n)
34+
if (colors[pre] == color) {
35+
remove(set, bit, n, pre)
36+
}
37+
val next = getIndex(i + 1, n)
38+
if (colors[next] == color) {
39+
remove(set, bit, n, i)
40+
}
41+
colors[i] = colors[i] xor 1
42+
color = colors[i]
43+
if (colors[pre] == color) {
44+
add(set, bit, n, pre)
45+
}
46+
if (colors[next] == color) {
47+
add(set, bit, n, i)
48+
}
4249
}
4350
}
51+
return ans
4452
}
4553

46-
// --- BIT wrapper methods ---
47-
// Updates both BITs for a given group length.
48-
private fun edt(x: Int, y: Int) {
49-
// Second BIT is updated with x * y.
50-
BITS[1].update(x, x * y)
51-
// First BIT is updated with y.
52-
BITS[0].update(x, y)
53-
}
54-
55-
// Combines BIT queries to get the result for a given x.
56-
private fun qry(x: Int): Int {
57-
return BITS[1].query(x) + (1 - x) * BITS[0].query(x)
58-
}
59-
60-
// Returns the length of a group from index x to y.
61-
private fun len(x: Int, y: Int): Int {
62-
return y - x + 1
63-
}
64-
65-
// --- Group operations ---
66-
// Removes a group (block) by updating BIT with a negative value.
67-
private fun removeGroup(start: Int, end: Int) {
68-
edt(len(start, end), -1)
69-
}
70-
71-
// Adds a group (block) by updating BIT with a positive value.
72-
private fun addGroup(start: Int, end: Int) {
73-
edt(len(start, end), 1)
74-
}
75-
76-
// Initializes the alternating groups using the colors array.
77-
private fun initializeGroups(colors: IntArray, groups: TreeMap<Int, Int>) {
78-
val n = colors.size
79-
var i = 0
80-
while (i < n) {
81-
var r = i
82-
// Determine the end of the current alternating group.
83-
while (r < n && (colors[r] + colors[i] + r + i) % 2 == 0) {
84-
++r
85-
}
86-
// The group spans from index i to r-1.
87-
groups.put(i, r - 1)
88-
// Update BITs with the length of this group.
89-
edt(r - i, 1)
90-
// Skip to the end of the current group.
91-
i = r - 1
92-
++i
54+
private fun add(set: BitSet, bit: BIT, n: Int, i: Int) {
55+
if (set.isEmpty) {
56+
bit.update(n, 1)
57+
} else {
58+
update(set, bit, n, i, 1)
9359
}
60+
set.set(i)
9461
}
9562

96-
// Processes a type 1 query: returns the number of alternating groups
97-
// of at least the given size.
98-
private fun processQueryType1(colors: IntArray, groups: TreeMap<Int, Int>, groupSize: Int): Int {
99-
var ans = qry(groupSize)
100-
val firstGroup = groups.firstEntry()
101-
val lastGroup = groups.lastEntry()
102-
// If there is more than one group and the first and last colors differ,
103-
// adjust the answer by "merging" the groups at the boundaries.
104-
if (firstGroup !== lastGroup && colors[0] != colors[colors.size - 1]) {
105-
val leftLen = len(firstGroup.key!!, firstGroup.value!!)
106-
val rightLen = len(lastGroup.key!!, lastGroup.value!!)
107-
ans = (ans - max((leftLen - groupSize + 1).toDouble(), 0.0)).toInt()
108-
ans = (ans - max((rightLen - groupSize + 1).toDouble(), 0.0)).toInt()
109-
ans = (ans + max((leftLen + rightLen - groupSize + 1).toDouble(), 0.0)).toInt()
63+
private fun remove(set: BitSet, bit: BIT, n: Int, i: Int) {
64+
set.clear(i)
65+
if (set.isEmpty) {
66+
bit.update(n, -1)
67+
} else {
68+
update(set, bit, n, i, -1)
11069
}
111-
return ans
11270
}
11371

114-
// Processes a type 2 query: updates the color at index x and adjusts groups.
115-
private fun processQueryType2(
116-
colors: IntArray,
117-
groups: TreeMap<Int, Int>,
118-
x: Int,
119-
newColor: Int,
120-
) {
121-
if (colors[x] == newColor) {
122-
return
72+
private fun update(set: BitSet, bit: BIT, n: Int, i: Int, v: Int) {
73+
var pre = set.previousSetBit(i)
74+
if (pre == -1) {
75+
pre = set.previousSetBit(n)
12376
}
124-
// Update the color at index x.
125-
colors[x] = newColor
126-
// Find the group (block) that contains index x.
127-
var it = groups.floorEntry(x)
128-
val l: Int = it!!.key!!
129-
val r: Int = it.value!!
130-
// Remove the old group from BIT and map.
131-
removeGroup(l, r)
132-
groups.remove(l)
133-
var ml = x
134-
var mr = x
135-
// Process the left side of index x.
136-
if (l != x) {
137-
groups.put(l, x - 1)
138-
addGroup(l, x - 1)
139-
} else {
140-
if (x > 0 && colors[x] != colors[x - 1]) {
141-
it = groups.floorEntry(x - 1)
142-
if (it != null) {
143-
ml = it.key!!
144-
removeGroup(it.key!!, it.value!!)
145-
groups.remove(it.key)
146-
}
147-
}
77+
var next = set.nextSetBit(i)
78+
if (next == -1) {
79+
next = set.nextSetBit(0)
14880
}
149-
// Process the right side of index x.
150-
if (r != x) {
151-
groups.put(x + 1, r)
152-
addGroup(x + 1, r)
153-
} else {
154-
if (x + 1 < colors.size && colors[x + 1] != colors[x]) {
155-
it = groups.ceilingEntry(x + 1)
156-
if (it != null) {
157-
mr = it.value!!
158-
removeGroup(it.key!!, it.value!!)
159-
groups.remove(it.key)
160-
}
161-
}
162-
}
163-
164-
// Merge the affected groups into one new group.
165-
groups.put(ml, mr)
166-
addGroup(ml, mr)
81+
bit.update(getIndex(next - pre + n - 1, n) + 1, -v)
82+
bit.update(getIndex(i - pre, n), v)
83+
bit.update(getIndex(next - i, n), v)
16784
}
16885

169-
// Clears both BITs. This is done after processing all queries.
170-
private fun clearAllBITs(n: Int) {
171-
for (i in 0..n + 2) {
172-
BITS[0].clear(i)
173-
BITS[1].clear(i)
174-
}
86+
private fun getIndex(index: Int, mod: Int): Int {
87+
val result = if (index >= mod) index - mod else index
88+
return if (index < 0) index + mod else result
17589
}
17690

177-
// Main function to handle queries on alternating groups.
178-
fun numberOfAlternatingGroups(colors: IntArray, queries: Array<IntArray>): MutableList<Int> {
179-
val groups = TreeMap<Int, Int>()
180-
val n = colors.size
181-
val results: MutableList<Int> = ArrayList<Int>()
182-
// Initialize alternating groups.
183-
initializeGroups(colors, groups)
184-
// Process each query.
185-
for (query in queries) {
186-
if (query[0] == 1) {
187-
// Type 1 query: count alternating groups of at least a given size.
188-
val groupSize = query[1]
189-
val ans = processQueryType1(colors, groups, groupSize)
190-
results.add(ans)
191-
} else {
192-
// Type 2 query: update the color at a given index.
193-
val index = query[1]
194-
val newColor = query[2]
195-
processQueryType2(colors, groups, index, newColor)
91+
private class BIT(n: Int) {
92+
var n: Int = n + 1
93+
var tree1: IntArray = IntArray(n + 1)
94+
var tree2: IntArray = IntArray(n + 1)
95+
96+
fun update(size: Int, v: Int) {
97+
var i = size
98+
while (i > 0) {
99+
tree1[i] += v
100+
tree2[i] += v * size
101+
i -= i and -i
196102
}
197103
}
198-
// Clear BITs after processing all queries.
199-
clearAllBITs(n)
200-
return results
201-
}
202104

203-
companion object {
204-
private const val SZ = 63333
205-
private const val OFFSET: Int = SZ - 10
206-
private val BITS = arrayOf<BIT>(BIT(), BIT())
105+
fun query(size: Int): IntArray {
106+
var count = 0
107+
var sum = 0
108+
var i = size
109+
while (i < n) {
110+
count += tree1[i]
111+
sum += tree2[i]
112+
i += i and -i
113+
}
114+
return intArrayOf(count, sum)
115+
}
207116
}
208117
}

0 commit comments

Comments
 (0)