結果

問題 No.956 Number of Unbalanced
コンテスト
ユーザー qwewe
提出日時 2025-05-14 12:53:37
言語 PyPy3
(7.3.15)
結果
TLE  
実行時間 -
コード長 1,933 bytes
コンパイル時間 163 ms
コンパイル使用メモリ 82,356 KB
実行使用メモリ 290,012 KB
最終ジャッジ日時 2025-05-14 12:55:21
合計ジャッジ時間 4,117 ms
ジャッジサーバーID
（参考情報） 		judge5 / judge4
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 6
other TLE * 1 -- * 23
権限があれば一括ダウンロードができます

ソースコード

diff # 

import sys
from bisect import bisect_left
from collections import defaultdict

def main():
    input = sys.stdin.read().split()
    n = int(input[0])
    A = list(map(int, input[1:n+1]))
    
    freq = defaultdict(int)
    for num in A:
        freq[num] += 1
    
    total = 0
    
    for x in freq:
        B = [1 if a == x else -1 for a in A]
        prefix = [0] * (n + 1)
        for i in range(n):
            prefix[i+1] = prefix[i] + B[i]
        
        # Compress prefix values
        sorted_prefix = sorted(set(prefix))
        rank = {v: i+1 for i, v in enumerate(sorted_prefix)}  # 1-based indexing
        size = len(sorted_prefix)
        
        class FenwickTree:
            def __init__(self, size):
                self.size = size
                self.tree = [0] * (self.size + 2)
            
            def update(self, idx):
                while idx <= self.size:
                    self.tree[idx] += 1
                    idx += idx & -idx
            
            def query(self, idx):
                res = 0
                while idx > 0:
                    res += self.tree[idx]
                    idx -= idx & -idx
                return res
        
        ft = FenwickTree(size)
        ans = 0
        # Initialize with prefix[0]
        s0 = prefix[0]
        idx = bisect_left(sorted_prefix, s0)
        if idx < len(sorted_prefix) and sorted_prefix[idx] == s0:
            ft.update(rank[s0])
        else:
            # This should not happen as sorted_prefix contains all elements
            pass
        
        for i in range(1, n+1):
            s = prefix[i]
            # Find the number of elements < s in the sorted_prefix
            cnt = bisect_left(sorted_prefix, s)
            ans += ft.query(cnt)
            # Update Fenwick Tree with current s
            ft.update(rank[s])
        
        total += ans
    
    print(total)

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