結果

問題 No.59 鉄道の旅
ユーザー lam6er
提出日時 2025-03-20 21:15:05
言語 PyPy3
(7.3.15)
結果
AC  
実行時間 203 ms / 5,000 ms
コード長 1,985 bytes
コンパイル時間 513 ms
コンパイル使用メモリ 82,848 KB
実行使用メモリ 118,396 KB
最終ジャッジ日時 2025-03-20 21:15:44
合計ジャッジ時間 2,374 ms
ジャッジサーバーID
（参考情報） 		judge1 / judge5
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 4
other AC * 12
権限があれば一括ダウンロードができます

ソースコード

diff # 

import bisect

class FenwickTree:
    def __init__(self, size):
        self.n = size
        self.tree = [0] * (self.n + 2)  # 1-based indexing

    def update(self, idx, delta):
        while idx <= self.n:
            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

def main():
    import sys
    input = sys.stdin.read().split()
    ptr = 0
    N = int(input[ptr])
    ptr += 1
    K = int(input[ptr])
    ptr += 1
    Ws = []
    all_ws = []
    for _ in range(N):
        w = int(input[ptr])
        ptr += 1
        Ws.append(w)
        all_ws.append(abs(w))  # for removal and addition

    # Create sorted unique list of W's
    sorted_unique_ws = sorted(list(set(all_ws)))
    M = len(sorted_unique_ws)

    # Initialize Fenwick Tree
    fenwick = FenwickTree(M)

    for w in Ws:
        if w > 0:
            # Add operation
            x = w
            # Find the index in sorted_unique_ws
            idx = bisect.bisect_left(sorted_unique_ws, x)
            if idx < len(sorted_unique_ws) and sorted_unique_ws[idx] == x:
                fenwick_idx = idx + 1  # 1-based
                # Calculate sum_ge = sum >= x
                sum_ge = fenwick.query(M) - fenwick.query(fenwick_idx - 1)
                if sum_ge < K:
                    fenwick.update(fenwick_idx, 1)
        else:
            # Remove operation
            x = -w
            idx = bisect.bisect_left(sorted_unique_ws, x)
            if idx < len(sorted_unique_ws) and sorted_unique_ws[idx] == x:
                fenwick_idx = idx + 1
                # Check current count
                current = fenwick.query(fenwick_idx) - fenwick.query(fenwick_idx - 1)
                if current > 0:
                    fenwick.update(fenwick_idx, -1)

    # Total sum is the answer
    print(fenwick.query(M))

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