Painting a Grid With Three Different Colors

Problem Description

So, you want to paint a grid with three different colors, huh? Sounds easy, right? Just grab a brush and go wild! But wait, there’s a catch! You can’t have the same color in adjacent cells, and you can’t have the same color in the same column. It’s like trying to pick an outfit for a party where your ex is also invited—awkward!

Imagine you have a grid of size m x n, and you need to fill it with colors represented by integers 1, 2, and 3. The goal is to find out how many ways you can paint this grid while following the rules.

Code Solution

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

        kMod = 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 %= kMod

            return ans

        return dp(0, 0, 0, 0)
    

Approach Explanation

The code uses a dynamic programming approach with memoization to efficiently calculate the number of valid ways to paint the grid. It defines two helper functions: getColor and setColor, which manage the colors in a bitmask format. The main function dp recursively explores all possible colorings while adhering to the constraints.

1. Base Cases: If the current column is filled, it returns 1. If the current row exceeds the grid height, it moves to the next column.
2. Color Selection: It iterates through the three colors, checking if the current color can be used based on the previous row and column constraints.
3. Recursion: It accumulates the results from valid configurations and returns the total count.

Time and Space Complexity

Time Complexity: O(m * n * 3^m), where m is the number of rows and n is the number of columns. The factor of 3m arises from the three color choices for each row.

Space Complexity: O(m * n), due to the memoization cache storing results for each state.

Real-World Example

Think of this problem like organizing a family reunion where you have three different colored shirts (red, green, blue) and you want to make sure no two family members wearing the same color stand next to each other. You want to maximize the number of unique arrangements while keeping the peace—because we all know how family gatherings can get!