Skip to content

Improved task 3245 #1943

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Mar 14, 2025
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
259 changes: 91 additions & 168 deletions src/main/java/g3201_3300/s3245_alternating_groups_iii/Solution.java
Original file line number Diff line number Diff line change
@@ -1,198 +1,121 @@
package g3201_3300.s3245_alternating_groups_iii;

// #Hard #Array #Binary_Indexed_Tree #2025_02_12_Time_135_ms_(86.36%)_Space_84.24_MB_(40.91%)
// #Hard #Array #Binary_Indexed_Tree #2025_03_14_Time_38_ms_(91.84%)_Space_77.53_MB_(87.76%)

import java.util.ArrayList;
import java.util.BitSet;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;

public class Solution {
private static final int SZ = 63333;
private static final int OFFSET = SZ - 10;
private static final BIT[] BITS = {new BIT(), new BIT()};

// Binary Indexed Tree (BIT) class.
private static class BIT {
int[] bs = new int[SZ];

// Update BIT: add value y to index x.
void update(int x, int y) {
x = OFFSET - x;
for (; x < SZ; x += x & -x) {
bs[x] += y;
}
}

// Query BIT: get the prefix sum up to index x.
int query(int x) {
x = OFFSET - x;
int ans = 0;
for (; x > 0; x -= x & -x) {
ans += bs[x];
public List<Integer> numberOfAlternatingGroups(int[] colors, int[][] queries) {
int n = colors.length;
BitSet set = new BitSet();
BIT bit = new BIT(n);
for (int i = 0; i < n; i++) {
if (colors[i] == colors[getIndex(i + 1, n)]) {
add(set, bit, n, i);
}
return ans;
}

// Clear BIT values starting from index x.
void clear(int x) {
x = OFFSET - x;
for (; x < SZ; x += x & -x) {
bs[x] = 0;
List<Integer> ans = new ArrayList<>();
for (int[] q : queries) {
if (q[0] == 1) {
if (set.isEmpty()) {
ans.add(n);
} else {
int size = q[1];
int[] res = bit.query(size);
ans.add(res[1] - res[0] * (size - 1));
}
} else {
int i = q[1];
int color = colors[i];
if (q[2] == color) {
continue;
}
int pre = getIndex(i - 1, n);
if (colors[pre] == color) {
remove(set, bit, n, pre);
}
int next = getIndex(i + 1, n);
if (colors[next] == color) {
remove(set, bit, n, i);
}
colors[i] ^= 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 void edt(int x, int y) {
// 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 int qry(int x) {
return BITS[1].query(x) + (1 - x) * BITS[0].query(x);
}

// Returns the length of a group from index x to y.
private int len(int x, int y) {
return y - x + 1;
}

// --- Group operations ---
// Removes a group (block) by updating BIT with a negative value.
private void removeGroup(int start, int end) {
edt(len(start, end), -1);
}

// Adds a group (block) by updating BIT with a positive value.
private void addGroup(int start, int end) {
edt(len(start, end), 1);
}

// Initializes the alternating groups using the colors array.
private void initializeGroups(int[] colors, TreeMap<Integer, Integer> groups) {
int n = colors.length;
int i = 0;
while (i < n) {
int 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 void add(BitSet set, BIT bit, int n, int i) {
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 int processQueryType1(int[] colors, TreeMap<Integer, Integer> groups, int groupSize) {
int ans = qry(groupSize);
Map.Entry<Integer, Integer> firstGroup = groups.firstEntry();
Map.Entry<Integer, Integer> 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.length - 1]) {
int leftLen = len(firstGroup.getKey(), firstGroup.getValue());
int rightLen = len(lastGroup.getKey(), lastGroup.getValue());
ans -= Math.max(leftLen - groupSize + 1, 0);
ans -= Math.max(rightLen - groupSize + 1, 0);
ans += Math.max(leftLen + rightLen - groupSize + 1, 0);
private void remove(BitSet set, BIT bit, int n, int i) {
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 void processQueryType2(
int[] colors, TreeMap<Integer, Integer> groups, int x, int newColor) {
if (colors[x] == newColor) {
return;
}
// Update the color at index x.
colors[x] = newColor;
// Find the group (block) that contains index x.
Map.Entry<Integer, Integer> it = groups.floorEntry(x);
int l = it.getKey();
int r = it.getValue();
// Remove the old group from BIT and map.
removeGroup(l, r);
groups.remove(l);
int ml = x;
int 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.getKey();
removeGroup(it.getKey(), it.getValue());
groups.remove(it.getKey());
}
}
private void update(BitSet set, BIT bit, int n, int i, int v) {
int pre = set.previousSetBit(i);
if (pre == -1) {
pre = set.previousSetBit(n);
}
// Process the right side of index x.
if (r != x) {
groups.put(x + 1, r);
addGroup(x + 1, r);
} else {
if (x + 1 < colors.length && colors[x + 1] != colors[x]) {
it = groups.ceilingEntry(x + 1);
if (it != null) {
mr = it.getValue();
removeGroup(it.getKey(), it.getValue());
groups.remove(it.getKey());
}
}
int next = set.nextSetBit(i);
if (next == -1) {
next = set.nextSetBit(0);
}
bit.update(getIndex(next - pre + n - 1, n) + 1, -v);
bit.update(getIndex(i - pre, n), v);
bit.update(getIndex(next - i, n), v);
}

// Merge the affected groups into one new group.
groups.put(ml, mr);
addGroup(ml, mr);
private int getIndex(int index, int mod) {
int result = index >= mod ? index - mod : index;
return index < 0 ? index + mod : result;
}

// Clears both BITs. This is done after processing all queries.
private void clearAllBITs(int n) {
for (int i = 0; i <= n + 2; ++i) {
BITS[0].clear(i);
BITS[1].clear(i);
private static class BIT {
int n;
int[] tree1;
int[] tree2;

BIT(int n) {
this.n = n + 1;
tree1 = new int[n + 1];
tree2 = new int[n + 1];
}
}

// Main function to handle queries on alternating groups.
public List<Integer> numberOfAlternatingGroups(int[] colors, int[][] queries) {
TreeMap<Integer, Integer> groups = new TreeMap<>();
int n = colors.length;
List<Integer> results = new ArrayList<>();
// Initialize alternating groups.
initializeGroups(colors, groups);
// Process each query.
for (int[] query : queries) {
if (query[0] == 1) {
// Type 1 query: count alternating groups of at least a given size.
int groupSize = query[1];
int ans = processQueryType1(colors, groups, groupSize);
results.add(ans);
} else {
// Type 2 query: update the color at a given index.
int index = query[1];
int newColor = query[2];
processQueryType2(colors, groups, index, newColor);
void update(int size, int v) {
for (int i = size; i > 0; i -= i & -i) {
tree1[i] += v;
tree2[i] += v * size;
}
}

int[] query(int size) {
int count = 0;
int sum = 0;
for (int i = size; i < n; i += i & -i) {
count += tree1[i];
sum += tree2[i];
}
return new int[] {count, sum};
}
// Clear BITs after processing all queries.
clearAllBITs(n);
return results;
}
}