結果

問題 No.1526 Sum of Mex 2
ユーザー gew1fw
提出日時 2025-06-12 19:59:56
言語 PyPy3
(7.3.15)
結果
TLE  
実行時間 -
コード長 2,370 bytes
コンパイル時間 143 ms
コンパイル使用メモリ 82,388 KB
実行使用メモリ 97,200 KB
最終ジャッジ日時 2025-06-12 20:03:13
合計ジャッジ時間 5,599 ms
ジャッジサーバーID
(参考情報)
judge3 / judge5
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
sample AC * 3
other AC * 10 TLE * 1 -- * 21
権限があれば一括ダウンロードができます

ソースコード

diff #

import sys
from collections import defaultdict

def main():
    N, *rest = map(int, sys.stdin.read().split())
    A = rest[:N]
    
    # Precompute the positions for each x
    pos = defaultdict(list)
    for i, a in enumerate(A):
        pos[a].append(i)
    
    total = 0
    
    # For x from 1 to N+1
    for x in range(1, N + 2):
        # Get the positions where x occurs, with sentinels
        x_pos = [-1] + [i for i, a in enumerate(A) if a == x] + [N]
        res = 0
        
        # Iterate through intervals between x's occurrences
        for i in range(1, len(x_pos)):
            L = x_pos[i-1] + 1
            R = x_pos[i] - 1
            if L > R:
                continue
            
            # Now, process the interval [L, R] which contains no x
            k = x - 1  # need to have all 1..k in the subarray
            if k == 0:
                # All subarrays in this interval contribute
                cnt = (R - L + 1) * (R - L + 2) // 2
                res += cnt
                continue
            
            # Check if all 1..k are present in the entire array
            # Precompute for efficiency
            present = [False] * (k + 1)
            for a in A:
                if 1 <= a <= k:
                    present[a] = True
            if not all(present[1:]):
                continue  # no subarrays can contribute
            
            # Now, compute the number of subarrays in [L, R] that contain all 1..k
            # Using sliding window
            freq = defaultdict(int)
            current = 0
            left = L
            cnt = 0
            for right in range(L, R + 1):
                a = A[right]
                if 1 <= a <= k:
                    if freq[a] == 0:
                        current += 1
                    freq[a] += 1
                # Shrink the window as much as possible
                while current == k and left <= right:
                    a_left = A[left]
                    if 1 <= a_left <= k:
                        freq[a_left] -= 1
                        if freq[a_left] == 0:
                            current -= 1
                    left += 1
                # The number of valid subarrays ending at right is left - L
                cnt += left - L
            res += cnt
        total += x * res
    print(total)

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