Lego Blocks

  • + 1 comment
    def legoBlocks(n, m):
        # If only one column or one row, there's only one way to build the wall
        if m == 1:
            return 1
        if n == 1:
            return 1 if 1 <= m <= 4 else 0
            
        MOD = pow(10, 9) + 7
        
        # Number of layouts to build a 1 x k wall for k = 1 to m
        p = [0, 1, 2, 4, 8]  # Base cases for widths 1 to 4
        for k in range(5, m + 1):
            # For width k, sum the previous four widths
            p.append(sum(p[-i] % MOD for i in range(1, 5)))
        
        # Number of layouts to build an n x k wall for k = 1 to m
        p = [pow(val, n, MOD) for val in p]
        
        # Number of good layouts (without any complete row gaps) to build n x k wall for k = 1 to m
        g = [0, 1]
        for k in range(2, m + 1):
            # Calculate the number of bad layouts for width k
            bad_layouts_num = sum(g[j] * p[k - j] for j in range(1, k)) % MOD
            # Explanation:
            # g[j] * p[k - j] represents the number of ways to build an n x k wall with a vertical division
            # at position j, resulting in a good layout on the left (g[j]) and any layout on the right (p[k - j]).
            # Summing these products for all j from 1 to k-1 gives the total number of bad layouts where there
            # is at least one complete vertical division in the wall.
            
            # Subtract bad layouts from total layouts to get good layouts
            g.append((p[k] - bad_layouts_num) % MOD)
        
        return g[m]