結果

問題 No.956 Number of Unbalanced
ユーザー lam6er
提出日時 2025-04-09 21:00:08
言語 PyPy3
(7.3.15)
結果
TLE  
実行時間 -
コード長 2,015 bytes
コンパイル時間 412 ms
コンパイル使用メモリ 81,984 KB
実行使用メモリ 292,948 KB
最終ジャッジ日時 2025-04-09 21:01:05
合計ジャッジ時間 4,292 ms
ジャッジサーバーID
（参考情報） 		judge1 / judge2
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 6
other TLE * 1 -- * 23
権限があれば一括ダウンロードができます

ソースコード

diff # 

import bisect
from collections import defaultdict

class FenwickTree:
    def __init__(self, size):
        self.n = size
        self.tree = [0] * (self.n + 1)
    
    def update(self, index, delta):
        while index <= self.n:
            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

def main():
    import sys
    input = sys.stdin.read().split()
    N = int(input[0])
    A = list(map(int, input[1:N+1]))
    
    # Create a frequency map of elements to their indices (1-based)
    pos_map = defaultdict(list)
    for idx, num in enumerate(A):
        pos_map[num].append(idx + 1)  # Store 1-based index
    
    total = 0
    
    for x in pos_map:
        positions = pos_map[x]
        # Compute prefix sums of x's occurrences up to each index (1-based)
        prefix = [0] * (N + 1)
        cnt = 0
        for i in range(1, N + 1):
            prefix[i] = cnt + (A[i-1] == x)
            cnt = prefix[i]
        
        # Calculate the S array
        S = [2 * prefix[i] - i for i in range(N + 1)]
        
        # Compress S values to handle large ranges
        sorted_S = sorted(S)
        unique_S = sorted(set(S))
        compressed = {v: i+1 for i, v in enumerate(unique_S)}
        
        # Initialize Fenwick Tree
        ft_size = len(unique_S)
        ft = FenwickTree(ft_size)
        
        res = 0
        for s in S:
            # Number of elements < current s is the leftmost position where sorted_S >= s
            # Using bisect_left gives the first index >= s, which is the count of elements < s.
            idx = bisect.bisect_left(unique_S, s)
            res += ft.query(idx)
            # Insert the compressed value of s into the Fenwick Tree
            c_idx = compressed[s]
            ft.update(c_idx, 1)
        total += res
    print(total)

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