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: