3245. Alternating Groups III

# Time:  O(nlogn + qlogn)
# Space: O(n)

from sortedcontainers import SortedList


# sorted list, freq table, bit, fenwick tree
class Solution(object):
    def numberOfAlternatingGroups(self, colors, queries):
        """
        :type colors: List[int]
        :type queries: List[List[int]]
        :rtype: List[int]
        """
        class BIT(object):  # 0-indexed.
            def __init__(self, n):
                self.__bit = [0]*(n+1)

            def add(self, i, val):
                i += 1
                while i < len(self.__bit):
                    self.__bit[i] += val
                    i += (i & -i)

            def query(self, i):
                i += 1
                ret = 0
                while i > 0:
                    ret += self.__bit[i]
                    i -= (i & -i)
                return ret

        def update(i, d):
            if d == +1:
                sl.add(i)
                if len(sl) == 1:
                    bit1.add(n, +1)
                    bit2.add(n, +n)
            curr = sl.index(i)
            prv, nxt = (curr-1)%len(sl), (curr+1)%len(sl)
            if len(sl) != 1:
                l = (sl[nxt]-sl[prv]-1)%n+1
                bit1.add(l, d*(-1))
                bit2.add(l, d*(-l))
                l = (sl[curr]-sl[prv])%n
                bit1.add(l, d*(+1))
                bit2.add(l, d*(+l))
                l = (sl[nxt]-sl[curr])%n
                bit1.add(l, d*(+1))
                bit2.add(l, d*(+l))
            if d == -1:
                if len(sl) == 1:
                    bit1.add(n, -1)
                    bit2.add(n, -n)
                sl.pop(curr)

        n = len(colors)
        sl = SortedList()
        bit1, bit2 = BIT(n+1), BIT(n+1)
        for i in xrange(n):
            if colors[i] == colors[(i+1)%n]:
                update(i, +1)
        result = []
        for q in queries:
            if q[0] == 1:
                l = q[1]
                result.append((bit2.query(n)-bit2.query(l-1))-
                              (l-1)*(bit1.query(n)-bit1.query(l-1)) if sl else n)
                continue
            _, i, c = q
            if colors[i] == c:
                continue      
            colors[i] = c          
            update((i-1)%n, +1 if colors[i] == colors[(i-1)%n] else -1) 
            update(i, +1 if colors[i] == colors[(i+1)%n] else -1)
        return result

Solution from kamyu104/LeetCode-Solutions · MIT