以下是 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)

 

Logo

中国智能体开发者社区,聚焦智能体与大模型开发,提供前沿资讯、实用工具链、开源项目及行业案例。通过技术沙龙、开发者大赛等活动,促进经验交流与协作,助力开发者快速构建创新智能应用。

更多推荐