結果

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

ソースコード

diff # 

import sys
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)
    pos = defaultdict(list)
    for idx, num in enumerate(a):
        freq[num] += 1
        pos[num].append(idx + 1)  # 1-based
    
    total = 0
    
    for x in pos:
        cnt = freq[x]
        if cnt == 1:
            total += 1
            continue
        
        # Compute prefix sums for transformed array
        prefix = [0] * (n + 1)
        for i in range(1, n + 1):
            if a[i-1] == x:
                prefix[i] = prefix[i-1] + 1
            else:
                prefix[i] = prefix[i-1] - 1
        
        # Collect all possible values for compression
        all_values = []
        for s in prefix:
            all_values.append(s)
            all_values.append(s - 1)
        
        # Compress the values
        sorted_values = sorted(list(set(all_values)))
        value_to_rank = {v: i+1 for i, v in enumerate(sorted_values)}  # 1-based
        
        # Fenwick Tree implementation
        class FenwickTree:
            def __init__(self, size):
                self.size = size
                self.tree = [0] * (self.size + 1)
            
            def update(self, idx, delta=1):
                while idx <= self.size:
                    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
        
        ft_size = len(sorted_values)
        fenwick = FenwickTree(ft_size)
        
        current_total = 0
        # Insert S[0]
        s0 = prefix[0]
        fenwick.update(value_to_rank[s0])
        
        for j in range(1, n + 1):
            s_j = prefix[j]
            target = s_j - 1
            
            # Find the largest value in sorted_values <= target
            left, right = 0, len(sorted_values) - 1
            best = -1
            while left <= right:
                mid = (left + right) // 2
                if sorted_values[mid] <= target:
                    best = mid
                    left = mid + 1
                else:
                    right = mid - 1
            
            if best == -1:
                cnt = 0
            else:
                rank = best + 1  # since sorted_values is 0-based, ranks are 1-based
                cnt = fenwick.query(rank)
            
            current_total += cnt
            
            # Add S[j] to Fenwick Tree (which is prefix[j] for the next iteration)
            s_prev = prefix[j]
            fenwick.update(value_to_rank[s_prev])
        
        total += current_total
    
    print(total)

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