結果

問題 No.956 Number of Unbalanced
ユーザー qwewe
提出日時 2025-04-24 12:32:10
言語 PyPy3
(7.3.15)
結果
TLE  
実行時間 -
コード長 1,614 bytes
コンパイル時間 209 ms
コンパイル使用メモリ 81,776 KB
実行使用メモリ 297,288 KB
最終ジャッジ日時 2025-04-24 12:33:24
合計ジャッジ時間 3,995 ms
ジャッジサーバーID
（参考情報） 		judge4 / judge1
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 6
other TLE * 1 -- * 23
権限があれば一括ダウンロードができます

ソースコード

diff # 

import bisect
from collections import defaultdict

def main():
    import sys
    input = sys.stdin.read().split()
    n = int(input[0])
    A = list(map(int, input[1:n+1]))
    
    pos_dict = defaultdict(list)
    for idx, num in enumerate(A):
        pos_dict[num].append(idx)
    
    total = 0
    
    for x in pos_dict:
        positions = pos_dict[x]
        k = len(positions)
        if k == 0:
            continue
        
        B = []
        for num in A:
            if num == x:
                B.append(1)
            else:
                B.append(-1)
        
        prefix = [0] * (n + 1)
        for i in range(n):
            prefix[i+1] = prefix[i] + B[i]
        
        s_list = list(prefix)
        sorted_s = sorted(set(s_list))
        rank = {v: i+1 for i, v in enumerate(sorted_s)}
        size = len(sorted_s)
        
        fenwick = [0] * (size + 2)
        
        def update(idx, val):
            while idx <= size:
                fenwick[idx] += val
                idx += idx & -idx
        
        def query(idx):
            res = 0
            while idx > 0:
                res += fenwick[idx]
                idx -= idx & -idx
            return res
        
        s0 = s_list[0]
        r0 = rank[s0]
        update(r0, 1)
        current = 0
        
        for i in range(1, n+1):
            s = prefix[i]
            pos = bisect.bisect_left(sorted_s, s)
            cnt = query(pos)
            current += cnt
            r = rank[s]
            update(r, 1)
        
        total += current
    
    print(total)

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