結果

問題 No.956 Number of Unbalanced
コンテスト
ユーザー gew1fw
提出日時 2025-06-12 15:24:12
言語 PyPy3
(7.3.15)
結果
TLE  
実行時間 -
コード長 1,595 bytes
コンパイル時間 174 ms
コンパイル使用メモリ 82,760 KB
実行使用メモリ 270,296 KB
最終ジャッジ日時 2025-06-12 15:24:25
合計ジャッジ時間 3,990 ms
ジャッジサーバーID
（参考情報） 		judge2 / judge5
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 6
other TLE * 1 -- * 23
権限があれば一括ダウンロードができます

ソースコード

diff # 

import sys
from collections import defaultdict

def main():
    sys.setrecursionlimit(1 << 25)
    n, *rest = list(map(int, sys.stdin.read().split()))
    A = rest[:n]
    
    class FenwickTree:
        def __init__(self, size):
            self.size = size
            self.tree = [0] * (self.size + 2)
        
        def update(self, index, delta=1):
            while index <= self.size:
                self.tree[index] += delta
                index += index & -index
        
        def query(self, index):
            res = 0
            while index > 0:
                res += self.tree[index]
                index -= index & -index
            return res
    
    max_A = max(A) if A else 0
    freq = defaultdict(list)
    for idx, num in enumerate(A):
        freq[num].append(idx)
    
    total = 0
    
    for x in freq:
        cnt = [0] * (n + 1)
        for i in range(1, n + 1):
            cnt[i] = cnt[i - 1] + (1 if A[i - 1] == x else 0)
        
        s = [0] * (n + 1)
        for i in range(n + 1):
            s[i] = 2 * cnt[i] - i
        
        min_s = min(s)
        max_s = max(s)
        ft_size = max_s - min_s + 2
        ft = FenwickTree(ft_size)
        ft.update(s[0] - min_s + 1)
        res = 0
        
        for j in range(1, n + 1):
            target = s[j] - 1 - min_s + 1
            if target <= 0:
                count = 0
            else:
                count = ft.query(target)
            res += count
            ft.update(s[j] - min_s + 1)
        
        total += res
    
    print(total)

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