結果

問題 No.833 かっこいい電車
ユーザー gew1fw
提出日時 2025-06-12 16:53:34
言語 PyPy3
(7.3.15)
結果
AC  
実行時間 542 ms / 2,000 ms
コード長 3,595 bytes
コンパイル時間 197 ms
コンパイル使用メモリ 82,552 KB
実行使用メモリ 119,368 KB
最終ジャッジ日時 2025-06-12 16:54:10
合計ジャッジ時間 9,496 ms
ジャッジサーバーID
(参考情報)
judge4 / judge3
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
sample AC * 2
other AC * 30
権限があれば一括ダウンロードができます

ソースコード

diff #

import sys

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

    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 range_query(self, l, r):
        if l > r:
            return 0
        return self.query(r) - self.query(l - 1)

class SegmentTree:
    def __init__(self, size):
        self.n = size
        if self.n == 0:
            self.size = 0
            self.tree = []
            return
        self.size = 1
        while self.size < self.n:
            self.size <<= 1
        self.tree = [0] * (2 * self.size)
    
    def update_point(self, pos, value):
        pos += self.size - 1  # convert to 0-based in the tree
        self.tree[pos] = value
        while pos > 1:
            pos >>= 1
            self.tree[pos] = self.tree[2 * pos] + self.tree[2 * pos + 1]
    
    def query_range(self, l, r):
        res = 0
        l += self.size - 1
        r += self.size - 1
        while l <= r:
            if l % 2 == 1:
                res += self.tree[l]
                l += 1
            if r % 2 == 0:
                res += self.tree[r]
                r -= 1
            l >>= 1
            r >>= 1
        return res

def find_l(x, seg_tree):
    low = 1
    high = x
    res = x
    while low <= high:
        mid = (low + high) // 2
        if mid > x - 1:
            res = mid
            high = mid - 1
            continue
        sum_edges = seg_tree.query_range(mid, x - 1)
        required = (x - 1 - mid + 1)
        if sum_edges == required:
            res = mid
            high = mid - 1
        else:
            low = mid + 1
    return res

def find_r(x, seg_tree, N):
    low = x
    high = N
    res = x
    while low <= high:
        mid = (low + high) // 2
        if mid == x:
            res = mid
            low = mid + 1
            continue
        if mid - 1 < x:
            res = mid
            low = mid + 1
            continue
        sum_edges = seg_tree.query_range(x, mid - 1)
        required = (mid - 1 - x + 1)
        if sum_edges == required:
            res = mid
            low = mid + 1
        else:
            high = mid - 1
    return res

def main():
    input = sys.stdin.read().split()
    ptr = 0
    N = int(input[ptr])
    ptr += 1
    Q = int(input[ptr])
    ptr += 1
    A = list(map(int, input[ptr:ptr+N]))
    ptr += N

    fenwick = FenwickTree(N)
    for i in range(N):
        fenwick.update(i + 1, A[i])

    if N == 1:
        seg_tree = None
    else:
        seg_tree = SegmentTree(N - 1)

    output = []
    for _ in range(Q):
        query_type = int(input[ptr])
        ptr += 1
        x = int(input[ptr])
        ptr += 1
        if query_type == 1:
            if N > 1:
                seg_tree.update_point(x, 1)
        elif query_type == 2:
            if N > 1:
                seg_tree.update_point(x, 0)
        elif query_type == 3:
            fenwick.update(x, 1)
        elif query_type == 4:
            if N == 1:
                l = 1
                r = 1
            else:
                l = find_l(x, seg_tree)
                r = find_r(x, seg_tree, N)
            total = fenwick.range_query(l, r)
            output.append(str(total))
    
    print('\n'.join(output))

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