結果

問題 No.121 傾向と対策：門松列（その２）
ユーザー lam6er
提出日時 2025-03-31 17:30:40
言語 PyPy3
(7.3.15)
結果
AC  
実行時間 4,508 ms / 5,000 ms
コード長 3,617 bytes
コンパイル時間 181 ms
コンパイル使用メモリ 82,056 KB
実行使用メモリ 462,492 KB
最終ジャッジ日時 2025-03-31 17:31:45
合計ジャッジ時間 15,433 ms
ジャッジサーバーID
（参考情報） 		judge3 / judge1
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
other AC * 9
権限があれば一括ダウンロードができます

ソースコード

diff # 

def main():
    import sys
    input = sys.stdin.read().split()
    idx = 0
    N = int(input[idx])
    idx += 1
    A = list(map(int, input[idx:idx + N]))
    idx += N

    # Coordinate compression
    unique = sorted(set(A))
    compress = {x: i + 1 for i, x in enumerate(unique)}  # 1-based index
    M = len(unique) + 2

    compressed_A = [compress[x] for x in A]

    # Precompute right_less and right_greater using Fenwick Tree
    right_less = [0] * N
    right_greater = [0] * N
    ft = FenwickTree(M)
    for j in range(N - 1, -1, -1):
        v = compressed_A[j]
        right_less[j] = ft.query(v - 1)
        right_greater[j] = ft.size() - ft.query(v)
        ft.add(v, 1)

    # Precompute left_less and left_greater using Fenwick Tree
    left_less = [0] * N
    left_greater = [0] * N
    ft_lt = FenwickTree(M)
    for j in range(N):
        v = compressed_A[j]
        left_less[j] = ft_lt.query(v - 1)
        cnt = j  # number of elements added so far (j elements when processing 0-based)
        left_greater[j] = cnt - ft_lt.query(v)
        ft_lt.add(v, 1)

    # Initialize right_count and left_count
    from collections import defaultdict
    right_count = defaultdict(int)
    for x in compressed_A:
        right_count[x] += 1
    left_count = defaultdict(int)

    # Fenwick Tree for same_less and same_greater
    class Fenwick:
        def __init__(self, size):
            self.n = size
            self.tree = [0] * (self.n + 2)

        def update(self, idx, delta):
            while idx <= self.n:
                self.tree[idx] += delta
                idx += idx & -idx

        def query(self, idx):
            res = 0
            while idx > 0:
                res += self.tree[idx]
                idx -= idx & -idx
            return res

        def query_range(self, a, b):
            if a > b:
                return 0
            return self.query(b) - self.query(a - 1)

    fen_less_greater = Fenwick(M)

    # Initialize Fenwick Tree with left[x] * right[x] for all x (initial left[x] is 0)
    # Initially, all products are 0

    ans = 0
    for j in range(N):
        v = compressed_A[j]

        # Decrement right_count[v]
        RC_old = right_count[v]
        right_count[v] -= 1

        # Calculate same_less and same_greater
        same_less = fen_less_greater.query(v - 1)
        same_greater = fen_less_greater.query_range(v + 1, M)

        # Case 1: A_j is max, so left_less * right_less - same_less
        case1 = left_less[j] * right_less[j] - same_less
        if case1 > 0:
            ans += case1

        # Case 2: A_j is min, left_greater * right_greater - same_greater
        case2 = left_greater[j] * right_greater[j] - same_greater
        if case2 > 0:
            ans += case2

        # Update left_count[v] and Fenwick tree
        LC_old = left_count[v]
        left_count[v] += 1

        # Calculate delta = (LC_old +1) * (RC_old -1) - (LC_old * RC_old)
        delta = ( (LC_old +1) * (RC_old -1) ) - (LC_old * RC_old)
        fen_less_greater.update(v, delta)

    print(ans)

class FenwickTree:
    def __init__(self, size):
        self.n = size
        self.tree = [0] * (self.n + 2)
        self.cnt = 0  # Number of elements in the tree

    def add(self, idx, delta):
        self.cnt += delta
        while idx <= self.n:
            self.tree[idx] += delta
            idx += idx & -idx

    def query(self, idx):
        res = 0
        while idx > 0:
            res += self.tree[idx]
            idx -= idx & -idx
        return res

    def size(self):
        return self.cnt

if __name__ == '__main__':
    main()
0