Grid Illumination LeetCode solution in C++

1
class Solution {
2
 public:
3
  // store info about rows , coulms and diagonal that is illuminated
4
  // if (i,j) is lamp , then every cell (x,y) is illuminated if
5
  // x==i (row) || y==j (column) || x+y == i+j (anti diagonal) || x-y == i-j
6
  // (main diagonal)
7
  vector<int> gridIllumination(int n, vector<vector<int>> &lamps, vector<vector<int>> &queries) {
8
    vector<int> ans;
9
    unordered_map<int, int> row, col, mainDiag, antiDiag;
10
    set<pair<int, int>> s;
11

12
    for (auto &l : lamps) {
13
      int x = l[0], y = l[1];
14
      // Even if the same lamp is listed more than once, it is turned on.
15
      // To prevent duplicate entries
16
      if (s.find({x, y}) == s.end()) {
17
        row[x]++;
18
        col[y]++;
19
        antiDiag[x + y]++;
20
        mainDiag[x - y]++;
21
        s.insert({x, y}); // This helps us to get the lamp in O(1)
22
      }
23
    }
24

25
    for (auto &q : queries) {
26
      // For every query {x,y} do 2 things :
27
      // 1. Check whether grid[x][y] is illuminated or not, accordingly set
28
      // answer either to '1' OR '0'
29
      // 2. if grid is illuminated, turn off the lamps at {x,y} and 8 adjacent
30
      // lamps = 9 lamps
31

32
      int x = q[0], y = q[1];
33
      if (row[x] > 0 or col[y] > 0 or antiDiag[x + y] > 0 or mainDiag[x - y] > 0) {
34
        // illumination
35
        ans.push_back(1);
36

37
        // Turn off the 9 lamps = 1 center and 8 adjacent lamps
38
        for (int i = -1; i <= 1; i++) {
39
          for (int j = -1; j <= 1; j++) {
40
            int nx = x + i;
41
            int ny = y + j;
42
            // NxN board, check boundaries condition
43
            if (nx >= 0 and nx < n and ny >= 0 and ny < n and s.find({nx, ny}) != s.end()) {
44
              row[nx]--;
45
              col[ny]--;
46
              antiDiag[nx + ny]--;
47
              mainDiag[nx - ny]--;
48
              s.erase({nx, ny});
49
            }
50
          }
51
        }
52
      } else // No illumination
53
        ans.push_back(0);
54
    }
55
    return ans;
56
  }
57
};