ユニークビジョンプログラミングコンテスト2024 春（AtCoder Beginner Contest 346）G問題メモ

from collections import defaultdict
from operator import add
from typing import TypeVar, Generic, Callable, List

T = TypeVar('T')
U = TypeVar('U')


class LazySegmentTreeInjectable(Generic[T, U]):
    def __init__(self,
                 n: int,
                 identity_data: Callable[[], T],
                 identity_lazy: Callable[[], U],
                 operation: Callable[[T, T], T],
                 mapping: Callable[[T, U], T],
                 composition: Callable[[U, U], U],
                 ):
        self.n = n
        self.depth = n.bit_length()
        self.offset = 1 << self.depth
        self.identity_data = identity_data
        self.identity_lazy = identity_lazy
        self.operation = operation
        self.mapping = mapping
        self.composition = composition
        self.data = [identity_data() for _ in range(self.offset << 1)]
        self.lazy = [identity_lazy() for _ in range(self.offset)]

    @classmethod
    def from_array(cls,
                   arr: List[T],
                   identity_data: Callable[[], T],
                   identity_lazy: Callable[[], U],
                   operation: Callable[[T, T], T],
                   mapping: Callable[[T, U], T],
                   composition: Callable[[U, U], U],
                   ):
        ins = cls(len(arr), identity_data, identity_lazy, operation, mapping, composition)
        ins.data[ins.offset:ins.offset + ins.n] = arr
        for i in range(ins.offset - 1, 0, -1):
            ins.update(i)
        return ins

    def init(self):
        for i in range(self.offset - 1, 0, -1):
            self.update(i)

    def push(self, i: int) -> None:
        if i < self.offset:
            data = self.data
            lazy = self.lazy
            lz = lazy[i]

            lch = i << 1
            rch = lch + 1
            data[lch] = self.mapping(data[lch], lz)
            data[rch] = self.mapping(data[rch], lz)
            if lch < self.offset:
                lazy[lch] = self.composition(lazy[lch], lz)
                lazy[rch] = self.composition(lazy[rch], lz)
            lazy[i] = self.identity_lazy()

    def update(self, i: int) -> None:
        lch = i << 1
        rch = lch + 1
        self.data[i] = self.operation(self.data[lch], self.data[rch])

    def all_apply(self, i: int, d: U) -> None:
        self.data[i] = self.mapping(self.data[i], d)
        if i < self.offset:
            self.lazy[i] = self.composition(self.lazy[i], d)

    def propagate(self, l: int, r: int) -> None:
        for i in range(self.depth, 0, -1):
            if ((l >> i) << i) != l:
                self.push(l >> i)
            if ((r >> i) << i) != r:
                self.push((r - 1) >> i)

    def range_update(self, l: int, r: int, d: U) -> None:
        l += self.offset
        r += self.offset
        self.propagate(l, r)

        l2 = l
        r2 = r
        while l < r:
            if (l & 1) == 1:
                self.all_apply(l, d)
                l += 1
            if (r & 1) == 1:
                r -= 1
                self.all_apply(r, d)
            l >>= 1
            r >>= 1
        l = l2
        r = r2

        for i in range(1, self.depth + 1):
            if ((l >> i) << i) != l:
                self.update(l >> i)
            if ((r >> i) << i) != r:
                self.update((r - 1) >> i)

    def range_query(self, l: int, r: int) -> T:
        l += self.offset
        r += self.offset

        self.propagate(l, r)
        sml = self.identity_data()
        smr = self.identity_data()
        while l < r:
            if (l & 1) == 1:
                sml = self.operation(sml, self.data[l])
                l += 1
            if (r & 1) == 1:
                r -= 1
                smr = self.operation(self.data[r], smr)
            l >>= 1
            r >>= 1
        return self.operation(sml, smr)

    def point_set(self, p: int, d: T) -> None:
        p += self.offset
        for i in range(self.depth, 0, -1):
            self.push(p >> i)
        self.data[p] = d
        for i in range(1, self.depth + 1):
            self.update(p >> i)

    def point_get(self, p: int) -> T:
        p += self.offset
        for i in range(self.depth, 0, -1):
            self.push(p >> i)
        return self.data[p]

    def debug_print(self) -> None:
        i = 1
        while i <= self.offset:
            # print(*map('{: 4d}'.format, self.data[i:i * 2]))
            print(*self.data[i:i * 2])
            i <<= 1
        i = 1
        while i <= self.offset:
            print(*map('{: 4d}'.format, self.lazy[i:i * 2]))
            # print(*self.lazy[i:i * 2])
            i <<= 1


def operation(a, b):
    if a[0] == b[0]:
        return (a[0], a[1] + b[1])
    if a[0] < b[0]:
        return a
    return b


def mapping(a, b):
    return (a[0] + b, a[1])


def solve(seg):
    n = len(seg)
    lzt = LazySegmentTreeInjectable.from_array([(0, 1) for _ in range(n)], lambda: (n, 0), int, operation, mapping, add)
    lzt.init()
    pos = defaultdict(list)
    for i, c in enumerate(seg):
        pos[c].append(i)
    progress = {}
    for c in pos:
        pos[c].append(n)
        pos[c].append(n)
        i1 = pos[c][0]
        i2 = pos[c][1]
        progress[c] = 1
        lzt.range_update(i1, i2, 1)
    res = 0
    for i in range(n):
        min_value, min_count = lzt.range_query(i, lzt.offset)
        if min_value == 0:
            res += n - i - min_count
        else:
            res += n - i
        c = seg[i]
        prg = progress[c]
        pc = pos[c]
        i0 = pc[prg - 1]
        i1 = pc[prg]
        i2 = pc[prg + 1]
        progress[c] += 1
        lzt.range_update(i0, i1, -1)
        lzt.range_update(i1, i2, 1)
    return res


n = int(input())
aaa = list(map(int, input().split()))

ans = solve(aaa)

print(ans)