時間計算量 \(O(N+Q)\) の解法における、簡潔な実装を紹介します。

ある部品 \(P\) が取り付けられるとき、「現在のロボットの重さ」は \(W_P\) だけ増加します。この後に、再び \(P\) のクエリが処理されるとき、「現在のロボットの重さ」は \(W_P\) だけ減少、つまり \(-W_P\) だけ増加します。(これは、部品が取り外されたときにも、同じようなことが言えます。)

よって、クエリが処理されたとき、\(W_P \leftarrow -W_P\) と更新することで、毎クエリの最初には \(W_P\) を加算するだけでよく、条件分岐が必要なくなります。

そのため、このように、陽にフラグ用の配列を持たずに実装することが可能です。

実装例(Piet)

実装例(C++)

#include <bits/stdc++.h>
using namespace std;
#define rep(i, n) for (int i = 0; i < (int)(n); ++i)

int main() {
    int X, N;
    cin >> X >> N;
    vector<int> W(N);
    rep(i, N) cin >> W[i];
    
    int Q;
    cin >> Q;
    rep(q, Q) {
        int P;
        cin >> P;
        --P;
        
        X += W[P];
        W[P] *= -1;
        
        cout << X << endl;
    }
}