Problem statement

https://leetcode.com/problems/matchsticks-to-square/

Solution

Special case of Problem 698 Partition to K Equal Sum Subsets, but here we need to partition into 4 equal sums, we say, put them into buckets (or sides of square).

Please also check my post https://leetcode.com/discuss/general-discussion/1125779/Dynamic-programming-on-subsets-with-examples-explained for more problems which can be solved with this technique.

Denote by dp(mask) possibility to put numbers from bitmask mask into say k subsets, such that sum inside each of the first k-1 subsets is equal to basket . In more detais: we return -1 if this partition is not possible and we return t >= 0, if it is possible, where t is remainder when we divide sum of all used numbers so far by basket: value, of how many we need to put in each group. Given our mask there is several options about what matchstick can be taken the last: the places, where we meet 1 bit. So, if neib is answer for smaller problem with this match removed dfs(mask ^ 1<<j), and neib >= 0, that is solution exists, and neib + nums[j] <= basket, that is we still have place in the last backet, then we return (neib + nums[j]) % basket. We use % here, because basked can be filled fully and in this case we need to start new one.

Complexity

We have O(2^n) states with O(n) transactions for each one, so overall time complexity is (2^n * n) and space complexity is O(2^n).

Code

class Solution:
    def makesquare(self, nums):
        N = len(nums)
        basket, rem = divmod(sum(nums), 4)
        if rem or max(nums) > basket: return False
        
        @lru_cache(None)
        def dfs(mask):
            if mask == 0: return 0
            for j in range(N):
                if mask & 1<<j:
                    neib = dfs(mask ^ 1<<j)
                    if neib >= 0 and neib + nums[j] <= basket:
                        return (neib + nums[j]) % basket
            return -1
                    
        return dfs((1<<N) - 1) == 0