以下是 LeetCode 3245. 交替组 III 的 TypeScript 实现。
核心思路与 Java 一致：Treap（树堆）维护相邻颜色相同的断点集合，树状数组（BIT）维护各交替段的长度，从而支持 O(\\log n) 的修改与查询。
—
完整代码
```typescript function numberOfAlternatingGroups(colors: number[], queries: number[][]): number[] { const n = colors.length; const set = new Treap(); // 维护断点位置（有序集合） const bit = new BIT(n); // 维护交替段长度的树状数组
// 计算环上从 a 到 b 的交替段长度 const getLen = (a: number, b: number): number => { let len = (b – a + n) % n; if (len === 0) len = n; return len; };
// 添加断点 i const add = (i: number): void => { if (set.isEmpty()) { bit.update(n, 1); // 整个环构成一个长度为 n 的段 } else { const pre = set.predecessor(i) ?? set.getMax()!; const nxt = set.successor(i) ?? set.getMin()!; const oldLen = getLen(pre, nxt); // 原来 pre→nxt 的整段被拆分 const len1 = getLen(pre, i); const len2 = getLen(i, nxt); bit.update(oldLen, -1); bit.update(len1, 1); bit.update(len2, 1); } set.insert(i); };
// 移除断点 i const remove = (i: number): void => { set.erase(i); if (set.isEmpty()) { bit.update(n, -1); } else { const pre = set.predecessor(i) ?? set.getMax()!; const nxt = set.successor(i) ?? set.getMin()!; const newLen = getLen(pre, nxt); // pre→nxt 两段合并 const len1 = getLen(pre, i); const len2 = getLen(i, nxt); bit.update(newLen, 1); bit.update(len1, -1); bit.update(len2, -1); } };
// 初始化：找出所有相邻颜色相同的断点 for (let i = 0; i < n; i++) { if (colors[i] === colors[(i + 1) % n]) { add(i); } }
const ans: number[] = []; for (const q of queries) { if (q[0] === 1) { // 查询：大小为 size 的交替组数量 const size = q[1]; if (set.isEmpty()) { ans.push(n); // 完全交替，任意 n 个起点都合法 } else { const [cnt, sum] = bit.query(size); // 每个长度为 L 的段贡献 L – size + 1 个组 ans.push(sum – cnt * (size – 1)); } } else { // 修改：colors[i] = newColor const i = q[1]; const newColor = q[2]; if (colors[i] === newColor) continue;
const pre = (i – 1 + n) % n; const nxt = (i + 1) % n;
// 修改前：移除旧断点 if (colors[pre] === colors[i]) remove(pre); if (colors[i] === colors[nxt]) remove(i);
colors[i] ^= 1; // 0↔1 翻转
// 修改后：添加新断点 if (colors[pre] === colors[i]) add(pre); if (colors[i] === colors[nxt]) add(i); } }
return ans; }
/* ==================== Treap（支持前驱/后继查询） ==================== */ class TreapNode { key: number; priority: number; left: TreapNode | null; right: TreapNode | null; constructor(key: number) { this.key = key; this.priority = Math.random(); this.left = null; this.right = null; } }
class Treap { root: TreapNode | null = null;
private rotateRight(p: TreapNode): TreapNode { const q = p.left!; p.left = q.right; q.right = p; return q; }
private rotateLeft(p: TreapNode): TreapNode { const q = p.right!; p.right = q.left; q.left = p; return q; }
insert(key: number): void { this.root = this._insert(this.root, key); }
private _insert(node: TreapNode | null, key: number): TreapNode { if (node === null) return new TreapNode(key); if (key === node.key) return node; if (key < node.key) { node.left = this._insert(node.left, key); if (node.left.priority > node.priority) node = this.rotateRight(node); } else { node.right = this._insert(node.right, key); if (node.right.priority > node.priority) node = this.rotateLeft(node); } return node; }
erase(key: number): void { this.root = this._erase(this.root, key); }
private _erase(node: TreapNode | null, key: number): TreapNode | null { if (node === null) return null; if (key === node.key) { if (node.left === null) return node.right; if (node.right === null) return node.left; if (node.left.priority > node.right.priority) { node = this.rotateRight(node); node.right = this._erase(node.right, key); } else { node = this.rotateLeft(node); node.left = this._erase(node.left, key); } } else if (key < node.key) { node.left = this._erase(node.left, key); } else { node.right = this._erase(node.right, key); } return node; }
// 严格小于 key 的最大值 predecessor(key: number): number | null { let res: number | null = null; let node = this.root; while (node !== null) { if (node.key < key) { res = node.key; node = node.right; } else { node = node.left; } } return res; }
// 严格大于 key 的最小值 successor(key: number): number | null { let res: number | null = null; let node = this.root; while (node !== null) { if (node.key > key) { res = node.key; node = node.left; } else { node = node.right; } } return res; }
getMin(): number | null { let node = this.root; if (!node) return null; while (node.left) node = node.left; return node.key; }
getMax(): number | null { let node = this.root; if (!node) return null; while (node.right) node = node.right; return node.key; }
isEmpty(): boolean { return this.root === null; } }
/* ==================== 树状数组（维护段长后缀信息） ==================== */ class BIT { n: number; tree1: number[]; // 维护个数 tree2: number[]; // 维护长度和
constructor(n: number) { this.n = n + 2; this.tree1 = new Array(n + 2).fill(0); this.tree2 = new Array(n + 2).fill(0); }
// 在位置 size 处增加 v update(size: number, v: number): void { for (let i = size; i > 0; i -= i & -i) { this.tree1[i] += v; this.tree2[i] += v * size; } }
// 查询长度 >= size 的段：返回 [个数, 长度总和] query(size: number): [number, number] { let count = 0, sum = 0; for (let i = size; i < this.n; i += i & -i) { count += this.tree1[i]; sum += this.tree2[i]; } return [count, sum]; } } ```
—
核心思路
结构 作用 Treap 存储所有相邻颜色相同的断点位置，支持 O(\\log n) 的插入、删除、前驱/后继查询（环状首尾相连） BIT 维护每个最大交替段的长度。`query(size)` 返回长度 \\ge size 的段的数量与长度之和
查询 `[1, size]`：
若环上无断点，答案为 `n`；否则利用公式
\\text{ans} = \\sum{\\text{段}}(L – size + 1) = \\left(\\sum L\\right) – \\text{count} \\times (size-1)
修改 `[2, i, color]`：
单点修改只影响与左右邻居 `(i-1, i)` 和 `(i, i+1)` 的相邻关系，最多触发 2 次断点删除 + 2 次断点添加，每次 O(\\log n)。
—
复杂度
– 时间：初始化 O(n \\log n)，每次查询/修改 O(\\log n) – 空间：O(n)
