Skip to content

3391. Design a 3D Binary Matrix with Efficient Layer Tracking 👍

  • Time:
    • Constructor: $O(n\log n)$
    • setCell(x: int, y: int, z: int): $O(\log n)$
    • unsetCell(x: int, y: int, z: int): $O(\log n)$
    • largestMatrix: $O(\log n)$
  • Space: $O(n + |\texttt{setCell(x: int, y: int, z: int)}|)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
class Matrix3D {
 public:
  Matrix3D(int n) {
    isSet = vector<vector<vector<bool>>>(
        n, vector<vector<bool>>(n, vector<bool>(n)));
    count = vector<int>(n);
    for (int x = 0; x < n; ++x)
      pairs.insert({0, x});
  }

  void setCell(int x, int y, int z) {
    if (isSet[x][y][z])
      return;
    isSet[x][y][z] = true;
    pairs.erase({count[x], x});
    pairs.insert({++count[x], x});
  }

  void unsetCell(int x, int y, int z) {
    if (!isSet[x][y][z])
      return;
    isSet[x][y][z] = false;
    pairs.erase({count[x], x});
    pairs.insert({--count[x], x});
  }

  int largestMatrix() {
    return pairs.rbegin()->second;
  }

 private:
  vector<vector<vector<bool>>> isSet;
  vector<int> count;  // count[x] := the number of cells set in the x-th layer

  struct PairCompare {
    bool operator()(const pair<int, int>& a, const pair<int, int>& b) const {
      return a.first == b.first ? b.second > a.second : b.first > a.first;
    }
  };

  set<pair<int, int>, PairCompare> pairs;  // (count[x], x)
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
class Matrix3D {
  public Matrix3D(int n) {
    isSet = new boolean[n][n][n];
    count = new int[n];
    for (int x = 0; x < n; ++x)
      pairs.add(new Pair<>(0, x));
  }

  public void setCell(int x, int y, int z) {
    if (isSet[x][y][z])
      return;
    isSet[x][y][z] = true;
    pairs.remove(new Pair<>(count[x], x));
    pairs.add(new Pair<>(++count[x], x));
  }

  public void unsetCell(int x, int y, int z) {
    if (!isSet[x][y][z])
      return;
    isSet[x][y][z] = false;
    pairs.remove(new Pair<>(count[x], x));
    pairs.add(new Pair<>(--count[x], x));
  }

  public int largestMatrix() {
    return pairs.first().getValue();
  }

  private boolean[][][] isSet;
  private int[] count; // count[x] := the number of cells set in the x-th layer
  private TreeSet<Pair<Integer, Integer>> pairs = new TreeSet<>( // (count[x], x)
      Comparator.comparing(Pair<Integer, Integer>::getKey, Comparator.reverseOrder())
          .thenComparing(Pair<Integer, Integer>::getValue, Comparator.reverseOrder()));
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
class Matrix3D:
  def __init__(self, n: int):
    self.isSet = set()
    # count[x] := the number of set cells in the x-th layer
    self.count = collections.Counter()
    # (count[x], x)
    self.pairs: SortedList = SortedList(key=lambda x: (-x[0], -x[1]))
    self.pairs.update((0, x) for x in range(n))

  def setCell(self, x: int, y: int, z: int) -> None:
    if (x, y, z) in self.isSet:
      return
    self.isSet.add((x, y, z))
    self.pairs.remove((self.count[x], x))
    self.count[x] += 1
    self.pairs.add((self.count[x], x))

  def unsetCell(self, x: int, y: int, z: int) -> None:
    if (x, y, z) not in self.isSet:
      return
    self.isSet.remove((x, y, z))
    self.pairs.remove((self.count[x], x))
    self.count[x] -= 1
    self.pairs.add((self.count[x], x))

  def largestMatrix(self) -> int:
    return self.pairs[0][1]