Check if an Original String Exists Given Two Encoded Strings LeetCode solution in Python

1
from functools import lru_cache
2

3

4
class Solution:
5
    def possiblyEquals(self, s1: str, s2: str) -> bool:
6

7
        def getValidPrefixLength(s, start):
8
            end = start
9
            while end < len(s) and s[end].isdigit():
10
                end += 1
11
            return end
12

13
        @lru_cache(None)
14
        def possibleLengths(s):
15
            """Return all possible lengths represented by numeric string s."""
16
            ans = {int(s)}
17
            for i in range(1, len(s)):
18
                # add all lengths by splitting numeric string s at i
19
                ans |= {
20
                    x + y
21
                    for x in possibleLengths(s[:i])
22
                    for y in possibleLengths(s[i:])
23
                }
24
            return ans
25

26
        @lru_cache(None)
27
        def dp(i, j, diff):
28
            """Return True if s1[i:] matches s2[j:] with given differences."""
29

30
            # If both have reached end return true if none of them are leading
31
            if i == len(s1) and j == len(s2):
32
                return diff == 0
33

34
            # s1 has not reached end and s1 starts with a digit
35
            if i < len(s1) and s1[i].isdigit():
36
                i2 = getValidPrefixLength(s1, i)
37
                for L in possibleLengths(s1[i:i2]):
38
                    # substract since lead of s2 decreases by L
39
                    if dp(i2, j, diff - L):
40
                        return True
41

42
            # s2 has not reached end and s2 starts with a digit
43
            elif j < len(s2) and s2[j].isdigit():
44
                j2 = getValidPrefixLength(s2, j)
45
                for L in possibleLengths(s2[j:j2]):
46
                    # add since lead of s2 increase by L
47
                    if dp(i, j2, diff + L):
48
                        return True
49

50
            # if none of them have integer prefix or a lead over the other
51
            elif diff == 0:
52
                # if only one of them has reached end or current alphabets are not the same
53
                if i == len(s1) or j == len(s2) or s1[i] != s2[j]:
54
                    return False
55
                # skip same alphabets
56
                return dp(i + 1, j + 1, 0)
57

58
            # if none of them have integer prefix & s2 lead over s1
59
            elif diff > 0:
60
                # no s1 to balance s2's lead
61
                if i == len(s1):
62
                    return False
63
                # move s1 pointer forward and reduce diff
64
                return dp(i + 1, j, diff - 1)
65

66
            # if none of them have integer prefix & s1 lead over s2
67
            else:
68
                # no s2 to balance s1's lead
69
                if j == len(s2):
70
                    return False
71
                # move s2 pointer forward and increase diff
72
                return dp(i, j + 1, diff + 1)
73

74
        # start with starts of both s1 and s2 with no lead by any of them
75
        return dp(0, 0, 0)