Problem statement

https://leetcode.com/problems/the-score-of-students-solving-math-expression/

Solution

Let dp(i, j) be all possible answers for string s[i:j+1]. Then we can calculate answers recursively and also use restriction that answers can be <= 1000: without it you will get TLE

Complexity

Time complexity is O(k^3 * 1000^2), because we have O(k^2) subproblems, for each of which you can have upto 1000 * 1000 * k attempts.

Code

class Solution:
    def scoreOfStudents(self, s, answers):
        @lru_cache(None)
        def dp(i, j):
            if i == j: return (int(s[i]),)
            ans = set()
            for k in range(i+1, j, 2):
                for s1, s2 in product(dp(i, k-1), dp(k+1, j)):
                    cand = s1 * s2 if s[k] == "*" else s1 + s2
                    if cand <= 1000: ans.add(cand) 
            return ans
        
        sols, corr = dp(0, len(s) - 1), 0
        for part in s.split("+"):
            corr += prod(int(i) for i in part.split("*"))
        
        ans = 0
        for x in answers:
            if x == corr: ans += 5
            elif x in sols: ans += 2
                
        return ans