結果

問題 No.1526 Sum of Mex 2
ユーザー lam6er
提出日時 2025-03-31 17:59:32
言語 PyPy3
(7.3.15)
結果
TLE  
実行時間 -
コード長 2,513 bytes
コンパイル時間 193 ms
コンパイル使用メモリ 82,112 KB
実行使用メモリ 110,228 KB
最終ジャッジ日時 2025-03-31 18:00:40
合計ジャッジ時間 8,872 ms
ジャッジサーバーID
(参考情報)
judge3 / judge5
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
sample AC * 3
other AC * 30 TLE * 1 -- * 1
権限があれば一括ダウンロードができます

ソースコード

diff #

def main():
    import sys
    input = sys.stdin.read().split()
    n = int(input[0])
    a = list(map(int, input[1:n+1]))
    
    # Precompute the positions for each number (1-based)
    from collections import defaultdict
    pos = defaultdict(list)
    for idx, num in enumerate(a):
        pos[num].append(idx)
    
    res = 0
    
    for x in range(0, n + 2):
        # Check if mex is x+1
        target = x + 1
        # Determine the segments where target does not appear
        split_points = [-1]
        if target in pos:
            split_points.extend(pos[target])
        split_points.append(n)
        # Iterate through each segment
        for i in range(len(split_points) - 1):
            left = split_points[i] + 1
            right = split_points[i + 1] - 1
            if left > right:
                continue
            # Check if x can be considered for mex x+1
            if x == 0:
                # mex=1 requires that segment contains no 1. But target is 1, and split_points is [positions of 1], so segments are without 1
                # mex=1 is the count of subarrays in this segment
                length = right - left + 1
                res += (length * (length + 1) // 2) * (x + 1)
                continue
            # For x >= 1, we need 1..x to be present in the segment
            # Check if all 1..x exist in the array
            exists = True
            for num in range(1, x + 1):
                if not pos[num]:
                    exists = False
                    break
            if not exists:
                continue
            # Check if in this segment, all 1..x are present
            # Precompute the positions within the segment
            last = [-1] * (x + 2)  # last occurrence of each number
            required = set(range(1, x + 1))
            current = set()
            left_min = -1
            cnt = 0
            l = left
            for r in range(left, right + 1):
                num_r = a[r]
                if 1 <= num_r <= x:
                    if last[num_r] < l:
                        current.add(num_r)
                    last[num_r] = r
                # Check if all required numbers are present
                if current == required:
                    # find the minimal last position
                    min_last = min(last[1:x + 1])
                    cnt += (min_last - l + 1)
                # else: can't contribute
            res += cnt * (x + 1)
    print(res)

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