BinarySearch 0920 Candy Race Taking Square Candies

Problem statement

https://binarysearch.com/problems/Candy-Race-Taking-Square-Candies/

Solution

Equal to Leetcode 1510. Stone Game IV.

Complexity

It is O(n sqrt(n)) for time and O(n) for space.

Code

class Solution:
    def solve(self, n):
        @lru_cache(None)
        def dfs(state):
            if state == 0: return False
            for i in range(1, int(math.sqrt(state))+1):
                if not dfs(state - i*i): return True
            return False
        return dfs(n)