Skip to content

1931. Painting a Grid With Three Different Colors 👍

  • Time: $O(n \cdot 3^m \cdot 2^m)$
  • Space: $O(n \cdot 4^m + 4^m \cdot 2^m)$
 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
43
44
45
46
47
48
49
50
51
52
class Solution {
 public:
  int colorTheGrid(int m, int n) {
    this->m = m;
    this->n = n;
    return dp(0, 0, 0, 0);
  }

 private:
  static constexpr int kMod = 1'000'000'007;
  int m;
  int n;
  vector<vector<int>> mem = vector<vector<int>>(1000, vector<int>(1024));

  int dp(int r, int c, int prevColMask, int currColMask) {
    if (c == n)
      return 1;
    if (mem[c][prevColMask])
      return mem[c][prevColMask];
    if (r == m)
      return dp(0, c + 1, currColMask, 0);

    int ans = 0;

    // 1 := red, 2 := green, 3 := blue
    for (int color = 1; color <= 3; ++color) {
      if (getColor(prevColMask, r) == color)
        continue;
      if (r > 0 && getColor(currColMask, r - 1) == color)
        continue;
      ans += dp(r + 1, c, prevColMask, setColor(currColMask, r, color));
      ans %= kMod;
    }

    if (r == 0)
      mem[c][prevColMask] = ans;

    return ans;
  }

  // e.g. __ __ __ __ __
  //      01 10 11 11 11
  //      R  G  B  B  B
  // getColor(0110111111, 3) -> G
  int getColor(int mask, int r) {
    return mask >> r * 2 & 3;
  }

  int setColor(int mask, int r, int color) {
    return mask | color << r * 2;
  }
};

 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
43
44
45
46
47
48
49
50
class Solution {
  public int colorTheGrid(int m, int n) {
    this.m = m;
    this.n = n;
    return dp(0, 0, 0, 0);
  }

  private static final int MOD = 1_000_000_007;
  private int m;
  private int n;
  private int[][] mem = new int[1000][1024];

  private int dp(int r, int c, int prevColMask, int currColMask) {
    if (c == n)
      return 1;
    if (mem[c][prevColMask] != 0)
      return mem[c][prevColMask];
    if (r == m)
      return dp(0, c + 1, currColMask, 0);

    int ans = 0;

    // 1 := red, 2 := green, 3 := blue
    for (int color = 1; color <= 3; ++color) {
      if (getColor(prevColMask, r) == color)
        continue;
      if (r > 0 && getColor(currColMask, r - 1) == color)
        continue;
      ans += dp(r + 1, c, prevColMask, setColor(currColMask, r, color));
      ans %= MOD;
    }

    if (r == 0)
      mem[c][prevColMask] = ans;

    return ans;
  }

  // e.g. __ __ __ __ __
  //      01 10 11 11 11
  //      R  G  B  B  B
  // getColor(0110111111, 3) -> G
  private int getColor(int mask, int r) {
    return mask >> r * 2 & 3;
  }

  private int setColor(int mask, int r, int color) {
    return mask | color << r * 2;
  }
}

 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
class Solution:
  def colorTheGrid(self, m: int, n: int) -> int:
    def getColor(mask: int, r: int) -> int:
      return mask >> r * 2 & 3

    def setColor(mask: int, r: int, color: int) -> int:
      return mask | color << r * 2

    MOD = 1_000_000_007

    @functools.lru_cache(None)
    def dp(r: int, c: int, prevColMask: int, currColMask: int) -> int:
      if c == n:
        return 1
      if r == m:
        return dp(0, c + 1, currColMask, 0)

      ans = 0

      # 1 := red, 2 := green, 3 := blue
      for color in range(1, 4):
        if getColor(prevColMask, r) == color:
          continue
        if r > 0 and getColor(currColMask, r - 1) == color:
          continue
        ans += dp(r + 1, c, prevColMask, setColor(currColMask, r, color))
        ans %= MOD

      return ans

    return dp(0, 0, 0, 0)