B - Robot Weight Editorial by MMNMM

より高速な解法

それぞれの種類のパーツが現在ついているかの情報とあわせて、「現在のロボットの重さ」を保存しておくことで、この問題をより高速に解くことができます。

それぞれのクエリでは、\(1\) つの部品がつけられるか外されるだけなので、ロボットの重さに \(1\) つの部品の重さを足すもしくは引くことで更新を完了することができます。

適切に実装することで時間計算量は \(O(N+Q)\) となり、\(N,Q\) が \(2\times10 ^ 5\) 程度になっても十分高速に解くことができます。

実装例は以下のようになります。

#include <iostream>
#include <vector>

int main() {
    using namespace std;
    int X, N;
    cin >> X >> N;
    vector<int> W(N);
    for (auto&& w : W)
        cin >> w;
    vector<bool> b(N); // b[i] := i 種類目のパーツがついていれば true, そうでなければ false
    int weight = X; // 現在の重さ

    int Q;
    cin >> Q;
    for (int q = 0; q < Q; ++q) {
        int P;
        cin >> P;
        --P;

        if (b[P]) { // すでについていたら
            b[P] = false; // 外して
            weight -= W[P]; // 重さも減らす
        } else { // ついていなければ
            b[P] = true; // つけて
            weight += W[P]; // 重さを増やす
        }

        cout << weight << endl;
    }
    return 0;
}
X = int(input())
N = int(input())
W = list(map(int, input().split()))

b = [False] * N # b[i] := i 種類目のパーツがついていれば True, そうでなければ False
weight = X # 現在の重さ

Q = int(input())

for q in range(Q):
    P = int(input())
    P -= 1

    if b[P]: # すでについていたら
        b[P] = False # 外して
        weight -= W[P] # 重さも減らす
    else: # ついていなければ
        b[P] = True # つけて
        weight += W[P] # 重さを増やす

    print(weight)

posted:
last update: