Official

D - Robot Weight Editorial by en_translator


If we can manage for each type of parts whether it is attached to the robot or not, we can process the queries by updating that information while computing the total weight based on that information.

To manage \(N\) flags of “attached/detached,” we can use language features like arrays or lists.

The following is sample code.

#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] := true if part i is attached; false otherwise

    int Q;
    cin >> Q;
    for (int q = 0; q < Q; ++q) {
        int P;
        cin >> P;
        --P;
        if (b[P]) { // If attached,
            b[P] = false; // detach it
        } else { // If detached
            b[P] = true; // attach it
        }

        // Find the weight
        int weight = X;
        for (int i = 0; i < N; ++i) {
            if (b[i]) { // If part i is attached,
                weight += W[i]; // add the weight
            }
        }
        cout << weight << endl;
        // Or one may `#include <numeric>` and write:
        // cout << transform_reduce(begin(W), end(W), begin(b), X) << endl;
    }
    return 0;
}
X = int(input())
N = int(input())
W = list(map(int, input().split()))

b = [False] * N # b[i] := True if part i is attached; False otherwise

Q = int(input())

for q in range(Q):
    P = int(input())
    P -= 1
    if b[P]: # If attached,
        b[P] = False # detach it
    else: # If detached,
        b[P] = True # attach it

    # Find the weight
    weight = X
    for i in range(N):
        if b[i]: # If part i is attaached,
            weight += W[i] # add the weight
    print(weight)
    # Or one may write:
    # print(sum((w if f else 0 for w, f in zip(W, b)), X))

posted:
last update: