Update square.cpp

1 files changed, 43 insertions, 52 deletions

diff --git a/square.cpp b/square.cpp
index ab83eec..7714056 100644
--- a/square.cpp
+++ b/square.cpp
@@ -1,84 +1,75 @@
 #include <bits/stdc++.h>
+#define f first
+#define s second
 using namespace std;
 typedef long long ll;
 
+// trick to iterate through directions
+constexpr int dx[4] = { 1, 0, -1, 0 }, dy[4] = { 0, 1, 0, -1 };
+
+int S[100001]; // skill levels
+vector<int> neighbor[100001]; // vector of neighbors for each competitor
+
 int main() {
 	if (fopen("in", "r")) freopen("in", "r", stdin), freopen("out", "w", stdout);
 
-	int T;
-	cin >> T;
-	for (int t = 0; t < T; ++t) {
-		cout << "Case #" << t + 1 << ": ";
+	int T; cin >> T;
+	for (int t = 1; t <= T; ++t) {
+		cout << "Case #" << t << ": ";
 
-		int R, C;
-		cin >> R >> C;
-		vector<vector<int>> S(R); // we have to use vector of vector because R or C could be 1e5
-		vector<pair<int, int>> check; // set of candidate competitors to check
+		int R, C; cin >> R >> C;
 		ll ans = 0, sum = 0;
+		vector<int> check; // set of candidate competitors to check
 		for (int i = 0; i < R; ++i) {
-			S[i].resize(C);
 			for (int j = 0; j < C; ++j) {
-				cin >> S[i][j];
-				sum += S[i][j];
-				check.emplace_back(i, j);
+				cin >> S[i * C + j]; sum += S[i * C + j];
+				check.push_back(i * C + j);
 			}
 		}
+
+		// checks if coordinates are out of range
+		auto valid = [&R, &C](const int& x, const int& y) {
+			return x >= 0 && x < R && y >= 0 && y < C;
+		};
 		
-		vector<vector<vector<int>>> neighbor(R);
+		// precompute neighbors
 		for (int i = 0; i < R; ++i) {
-			neighbor[i].resize(C);
 			for (int j = 0; j < C; ++j) {
-				neighbor[i][j].resize(4);
-				neighbor[i][j][0] = i > 0 ? i - 1 : -1;
-				neighbor[i][j][1] = i < R - 1 ? i + 1 : -1;
-				neighbor[i][j][2] = j > 0 ? j - 1 : -1;
-				neighbor[i][j][3] = j < C - 1 ? j + 1 : -1;
+				neighbor[i * C + j].resize(4);
+				for (int d = 0; d < 4; ++d) {
+					int x = i + dx[d], y = j + dy[d]; // get neighbor
+					neighbor[i * C + j][d] = valid(x, y) ? x * C + y : -1;
+				}
 			}
 		}
 
-		while (1) {
-			vector<pair<int, int>> elim; // vector to store eliminated competitors
-			for (auto& x : check) { // examine every candidate competitor
-				int i = x.first, j = x.second, sum = 0, num = 0;
-				if (S[i][j]) {
-					if (neighbor[i][j][0] != -1) sum += S[neighbor[i][j][0]][j], ++num;
-					if (neighbor[i][j][1] != -1) sum += S[neighbor[i][j][1]][j], ++num;
-					if (neighbor[i][j][2] != -1) sum += S[i][neighbor[i][j][2]], ++num;
-					if (neighbor[i][j][3] != -1) sum += S[i][neighbor[i][j][3]], ++num;
-
-					if (num * S[i][j] < sum) elim.emplace_back(i, j);
+		while (1) { // simulate the competition
+			vector<int> elim; // vector to store eliminated competitors
+			for (auto& c : check) { // examine every candidate competitor
+				int i = c / C, j = c % C, sum = 0, num = 0;
+				if (S[c]) {
+					for (int d = 0; d < 4; ++d) {
+						if (neighbor[c][d] != -1) sum += S[neighbor[c][d]], ++num;
+					}
+					if (num * S[c] < sum) elim.push_back(c);
 				}
 			}
 
 			check.clear(); // clear check for finding candidates for next round
 			ans += sum;
 
-			for (auto& x : elim) {
-				int i = x.first, j = x.second;
-				if (S[i][j]) {
-					sum -= S[i][j];
-					S[i][j] = 0;
-
-					if (neighbor[i][j][0] != -1) {
-						neighbor[neighbor[i][j][0]][j][1] = neighbor[i][j][1];
-						check.emplace_back(neighbor[i][j][0], j);
-					}
-					if (neighbor[i][j][1] != -1) {
-						neighbor[neighbor[i][j][1]][j][0] = neighbor[i][j][0];
-						check.emplace_back(neighbor[i][j][1], j);
-					}
-					if (neighbor[i][j][2] != -1) {
-						neighbor[i][neighbor[i][j][2]][3] = neighbor[i][j][3];
-						check.emplace_back(i, neighbor[i][j][2]);
-					}
-					if (neighbor[i][j][3] != -1) {
-						neighbor[i][neighbor[i][j][3]][2] = neighbor[i][j][2];
-						check.emplace_back(i, neighbor[i][j][3]);
+			for (auto& c : elim) {
+				int i = c / C, j = c % C;
+				if (S[c]) { // remove each competitor
+					sum -= S[c], S[c] = 0;
+					for (int d = 0; d < 4; ++d) if (neighbor[c][d] != -1) {
+						neighbor[neighbor[c][d]][(d + 2) % 4] = neighbor[c][(d + 2) % 4];
+						check.push_back(neighbor[c][d]);
 					}
 				}
 			}
 
-			if (elim.empty()) break;
+			if (elim.empty()) break; // if no one eliminated we can end the competition
 		}
 
 		cout << ans << '\n';