Skip to content

Commit 85b7dbb

Browse files
javadevjavadev
committed
Improved task 3245
1 parent 016779e commit 85b7dbb

File tree

1 file changed

+132
-96
lines changed
  • src/main/java/g3201_3300/s3245_alternating_groups_iii

1 file changed

+132
-96
lines changed

‎src/main/java/g3201_3300/s3245_alternating_groups_iii/Solution.java

+132-96
Original file line numberDiff line numberDiff line change
@@ -12,19 +12,19 @@ public class Solution {
1212
private static final int OFFSET = SZ - 10;
1313
private static final BIT[] BITS = {new BIT(), new BIT()};
1414

15-
// Class to represent the Binary Indexed Tree (BIT)
15+
// Binary Indexed Tree (BIT) class.
1616
private static class BIT {
1717
int[] bs = new int[SZ];
1818

19-
// Update BIT: add value y to index x
19+
// Update BIT: add value y to index x.
2020
void update(int x, int y) {
2121
x = OFFSET - x;
2222
for (; x < SZ; x += x & -x) {
2323
bs[x] += y;
2424
}
2525
}
2626

27-
// Query BIT: get the prefix sum up to index x
27+
// Query BIT: get the prefix sum up to index x.
2828
int query(int x) {
2929
x = OFFSET - x;
3030
int ans = 0;
@@ -34,7 +34,7 @@ int query(int x) {
3434
return ans;
3535
}
3636

37-
// Clear BIT values up to index x
37+
// Clear BIT values starting from index x.
3838
void clear(int x) {
3939
x = OFFSET - x;
4040
for (; x < SZ; x += x & -x) {
@@ -43,126 +43,162 @@ void clear(int x) {
4343
}
4444
}
4545

46-
// Wrapper functions for updating and querying the BITs
46+
// --- BIT wrapper methods ---
47+
// Updates both BITs for a given group length.
4748
private void edt(int x, int y) {
48-
// Update second BIT with product of index and value
49+
// Second BIT is updated with x * y.
4950
BITS[1].update(x, x * y);
50-
// Update first BIT with value
51+
// First BIT is updated with y.
5152
BITS[0].update(x, y);
5253
}
5354

55+
// Combines BIT queries to get the result for a given x.
5456
private int qry(int x) {
55-
// Query BITs and combine results
5657
return BITS[1].query(x) + (1 - x) * BITS[0].query(x);
5758
}
5859

59-
// Function to calculate the length between two indices
60+
// Returns the length of a group from index x to y.
6061
private int len(int x, int y) {
6162
return y - x + 1;
6263
}
6364

64-
// Main function to handle the queries
65-
public List<Integer> numberOfAlternatingGroups(int[] colors, int[][] queries) {
66-
// Map to store start and end indices of alternating groups
67-
TreeMap<Integer, Integer> c = new TreeMap<>();
65+
// --- Group operations ---
66+
// Removes a group (block) by updating BIT with a negative value.
67+
private void removeGroup(int start, int end) {
68+
edt(len(start, end), -1);
69+
}
70+
71+
// Adds a group (block) by updating BIT with a positive value.
72+
private void addGroup(int start, int end) {
73+
edt(len(start, end), 1);
74+
}
75+
76+
// Initializes the alternating groups using the colors array.
77+
private void initializeGroups(int[] colors, TreeMap<Integer, Integer> groups) {
6878
int n = colors.length;
69-
// Initialize alternating groups
70-
for (int i = 0; i < colors.length; ++i) {
79+
int i = 0;
80+
while (i < n) {
7181
int r = i;
72-
// Find end of the current alternating group
73-
while (r < colors.length && (colors[r] + colors[i] + r + i) % 2 == 0) {
82+
// Determine the end of the current alternating group.
83+
while (r < n && (colors[r] + colors[i] + r + i) % 2 == 0) {
7484
++r;
7585
}
76-
// Store group boundaries in map
77-
c.put(i, r - 1);
78-
// Update BITs with new group
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.
7989
edt(r - i, 1);
80-
// Move to the end of the current group
90+
// Skip to the end of the current group.
8191
i = r - 1;
92+
++i;
8293
}
83-
// List to store results for type 1 queries
84-
List<Integer> results = new ArrayList<>();
85-
// Process each query
86-
for (int[] q : queries) {
87-
if (q[0] == 1) {
88-
// Query type 1: Count alternating groups of a given size
89-
int ans = qry(q[1]);
90-
Map.Entry<Integer, Integer> a = c.firstEntry();
91-
Map.Entry<Integer, Integer> b = c.lastEntry();
92-
if (a != b) {
93-
// Check if merging groups is possible
94-
if (colors[0] != colors[colors.length - 1]) {
95-
int l1 = len(a.getKey(), a.getValue());
96-
int l2 = len(b.getKey(), b.getValue());
97-
// Subtract groups that are too small
98-
ans -= Math.max(l1 - q[1] + 1, 0);
99-
ans -= Math.max(l2 - q[1] + 1, 0);
100-
// Add merged group size
101-
ans += Math.max(l1 + l2 - q[1] + 1, 0);
102-
}
103-
} else if (colors[0] != colors[colors.length - 1]) {
104-
// If there's only one group, check if it can span the entire array
105-
ans = n;
106-
}
107-
// Store result for type 1 query
108-
results.add(ans);
109-
} else {
110-
// Query type 2: Update color at a given index
111-
int x = q[1];
112-
int y = q[2];
113-
if (colors[x] == y) {
114-
// If color is already correct, skip update
115-
continue;
116-
}
117-
// Update color
118-
colors[x] = y;
119-
// Find the block containing index x
120-
Map.Entry<Integer, Integer> it = c.floorEntry(x);
121-
assert it != null && it.getKey() <= x && it.getValue() >= x;
122-
int l = it.getKey();
123-
int r = it.getValue();
124-
// Remove the old block
125-
edt(len(it.getKey(), it.getValue()), -1);
126-
c.remove(it.getKey());
127-
int ml = x;
128-
int mr = x;
129-
// Update or split the affected blocks
130-
if (l != ml) {
131-
c.put(l, x - 1);
132-
edt(len(l, x - 1), 1);
133-
} else {
134-
if (x > 0 && colors[x] != colors[x - 1]) {
135-
it = c.floorEntry(x - 1);
136-
if (it != null) {
137-
ml = it.getKey();
138-
edt(len(it.getKey(), it.getValue()), -1);
139-
c.remove(it.getKey());
140-
}
141-
}
94+
}
95+
96+
// Processes a type 1 query: returns the number of alternating groups
97+
// of at least the given size.
98+
private int processQueryType1(int[] colors, TreeMap<Integer, Integer> groups, int groupSize) {
99+
int ans = qry(groupSize);
100+
Map.Entry<Integer, Integer> firstGroup = groups.firstEntry();
101+
Map.Entry<Integer, Integer> 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) {
105+
if (colors[0] != colors[colors.length - 1]) {
106+
int leftLen = len(firstGroup.getKey(), firstGroup.getValue());
107+
int rightLen = len(lastGroup.getKey(), lastGroup.getValue());
108+
ans -= Math.max(leftLen - groupSize + 1, 0);
109+
ans -= Math.max(rightLen - groupSize + 1, 0);
110+
ans += Math.max(leftLen + rightLen - groupSize + 1, 0);
111+
}
112+
} else if (colors[0] != colors[colors.length - 1]) {
113+
// In the special case when there's a single group but the
114+
// first and last colors differ, the whole array is counted.
115+
ans = colors.length;
116+
}
117+
return ans;
118+
}
119+
120+
// Processes a type 2 query: updates the color at index x and adjusts groups.
121+
private void processQueryType2(
122+
int[] colors, TreeMap<Integer, Integer> groups, int x, int newColor) {
123+
if (colors[x] == newColor) {
124+
return;
125+
}
126+
// Update the color at index x.
127+
colors[x] = newColor;
128+
// Find the group (block) that contains index x.
129+
Map.Entry<Integer, Integer> it = groups.floorEntry(x);
130+
int l = it.getKey();
131+
int r = it.getValue();
132+
// Remove the old group from BIT and map.
133+
removeGroup(l, r);
134+
groups.remove(l);
135+
int ml = x;
136+
int mr = x;
137+
// Process the left side of index x.
138+
if (l != x) {
139+
groups.put(l, x - 1);
140+
addGroup(l, x - 1);
141+
} else {
142+
if (x > 0 && colors[x] != colors[x - 1]) {
143+
it = groups.floorEntry(x - 1);
144+
if (it != null) {
145+
ml = it.getKey();
146+
removeGroup(it.getKey(), it.getValue());
147+
groups.remove(it.getKey());
142148
}
143-
if (r != mr) {
144-
c.put(x + 1, r);
145-
edt(len(x + 1, r), 1);
146-
} else {
147-
if (x + 1 < colors.length && colors[x + 1] != colors[x]) {
148-
it = c.ceilingEntry(x + 1);
149-
if (it != null) {
150-
mr = it.getValue();
151-
edt(len(it.getKey(), it.getValue()), -1);
152-
c.remove(it.getKey());
153-
}
154-
}
149+
}
150+
}
151+
// Process the right side of index x.
152+
if (r != x) {
153+
groups.put(x + 1, r);
154+
addGroup(x + 1, r);
155+
} else {
156+
if (x + 1 < colors.length && colors[x + 1] != colors[x]) {
157+
it = groups.ceilingEntry(x + 1);
158+
if (it != null) {
159+
mr = it.getValue();
160+
removeGroup(it.getKey(), it.getValue());
161+
groups.remove(it.getKey());
155162
}
156-
c.put(ml, mr);
157-
// Add new or modified block
158-
edt(len(ml, mr), 1);
159163
}
160164
}
161-
// Clear BITs after processing all queries
165+
166+
// Merge the affected groups into one new group.
167+
groups.put(ml, mr);
168+
addGroup(ml, mr);
169+
}
170+
171+
// Clears both BITs. This is done after processing all queries.
172+
private void clearAllBITs(int n) {
162173
for (int i = 0; i <= n + 2; ++i) {
163174
BITS[0].clear(i);
164175
BITS[1].clear(i);
165176
}
177+
}
178+
179+
// Main function to handle queries on alternating groups.
180+
public List<Integer> numberOfAlternatingGroups(int[] colors, int[][] queries) {
181+
TreeMap<Integer, Integer> groups = new TreeMap<>();
182+
int n = colors.length;
183+
List<Integer> results = new ArrayList<>();
184+
// Initialize alternating groups.
185+
initializeGroups(colors, groups);
186+
// Process each query.
187+
for (int[] query : queries) {
188+
if (query[0] == 1) {
189+
// Type 1 query: count alternating groups of at least a given size.
190+
int groupSize = query[1];
191+
int ans = processQueryType1(colors, groups, groupSize);
192+
results.add(ans);
193+
} else {
194+
// Type 2 query: update the color at a given index.
195+
int index = query[1];
196+
int newColor = query[2];
197+
processQueryType2(colors, groups, index, newColor);
198+
}
199+
}
200+
// Clear BITs after processing all queries.
201+
clearAllBITs(n);
166202
return results;
167203
}
168204
}

0 commit comments

Comments
 (0)