Skip to content

1397. Find All Good Strings

  • Time: $O(n \cdot |\texttt{evil}| \cdot 26) = O(n \cdot |\texttt{evil}|)$
  • Space: $O(n \cdot |\texttt{evil}| \cdot 2^2 + |\texttt{evil}| \cdot 26) = O(|\texttt{evil}| \cdot n)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
class Solution {
 public:
  int findGoodStrings(int n, string s1, string s2, string evil) {
    // dp[i][j][k1][k2] := # of good strings for s[i:] and there're already j
    // Matches with `evil`, where k1 is the 0/1 tight constraint for s1 and k2
    // Is the 0/1 tight constraint for s2
    dp.resize(n,
              vector<vector<vector<int>>>(
                  evil.length(), vector<vector<int>>(2, vector<int>(2, -1))));
    // nextMatchedCount[i][j] := # next matched evil count given that there're
    // Already i matches with `evil` and the current char is ('a' + j)
    nextMatchedCount.resize(evil.length(), vector<int>(26, -1));
    return find(s1, s2, evil, 0, 0, true, true, getLPS(evil));
  }

 private:
  constexpr static int kMod = 1'000'000'007;
  vector<vector<vector<vector<int>>>> dp;
  vector<vector<int>> nextMatchedCount;

  int find(const string& s1, const string& s2, const string& evil, int i,
           int matchedEvilCount, bool isS1Prefix, bool isS2Prefix,
           const vector<int>& evilLPS) {
    // s[:i] contains `evil`, so don't consider any ongoing strings
    if (matchedEvilCount == evil.length())
      return 0;
    // Run out of string, contributes one
    if (i == s1.length())
      return 1;
    int& ans = dp[i][matchedEvilCount][isS1Prefix][isS2Prefix];
    if (ans != -1)
      return ans;
    ans = 0;
    const char minChar = isS1Prefix ? s1[i] : 'a';
    const char maxChar = isS2Prefix ? s2[i] : 'z';
    for (char c = minChar; c <= maxChar; ++c) {
      const int nextMatchedEvilCount =
          getNextMatchedEvilCount(evil, matchedEvilCount, c, evilLPS);
      ans += find(s1, s2, evil, i + 1, nextMatchedEvilCount,
                  isS1Prefix && c == s1[i], isS2Prefix && c == s2[i], evilLPS);
      ans %= kMod;
    }
    return ans;
  }

  // Get Longest Prefix also Suffix
  vector<int> getLPS(const string& s) {
    vector<int> lps(s.length());
    for (int i = 1, j = 0; i < s.length(); ++i) {
      while (j > 0 && s[j] != s[i])
        j = lps[j - 1];
      if (s[i] == s[j])
        lps[i] = ++j;
    }
    return lps;
  }

  // J := the next index we're trying to match with `currChar`
  int getNextMatchedEvilCount(const string& evil, int j, char currChar,
                              const vector<int>& evilLPS) {
    int& ans = nextMatchedCount[j][currChar - 'a'];
    if (ans != -1)
      return ans;
    while (j > 0 && evil[j] != currChar)
      j = evilLPS[j - 1];
    return ans = (evil[j] == currChar ? j + 1 : j);
  }
};
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
class Solution {
  public int findGoodStrings(int n, String s1, String s2, String evil) {
    // dp[i][j][k1][k2] := # of good strings for s[i:] and there're already j
    // Matches with `evil`, where k1 is the 0/1 tight constraint for s1 and k2
    // Is the 0/1 tight constraint for s2
    dp = new Integer[n][evil.length()][2][2];
    // nextMatchedCount[i][j] := # next matched evil count given that there're
    // Already i matches with `evil` and the current char is ('a' + j)
    nextMatchedCount = new Integer[evil.length()][26];
    return find(s1, s2, evil, 0, 0, true, true, getLPS(evil));
  }

  private static final int kMod = 1_000_000_007;
  private Integer[][][][] dp;
  private Integer[][] nextMatchedCount;

  private int find(final String s1, final String s2, final String evil, int i, int matchedEvilCount,
                   boolean isS1Prefix, boolean isS2Prefix, int[] evilLPS) {
    // s[:i] contains `evil`, so don't consider any ongoing strings
    if (matchedEvilCount == evil.length())
      return 0;
    // Run out of string, contributes one
    if (i == s1.length())
      return 1;
    final int k1 = isS1Prefix ? 1 : 0;
    final int k2 = isS2Prefix ? 1 : 0;
    if (dp[i][matchedEvilCount][k1][k2] != null)
      return dp[i][matchedEvilCount][k1][k2];
    dp[i][matchedEvilCount][k1][k2] = 0;
    final char minChar = isS1Prefix ? s1.charAt(i) : 'a';
    final char maxChar = isS2Prefix ? s2.charAt(i) : 'z';
    for (char c = minChar; c <= maxChar; ++c) {
      final int nextMatchedEvilCount = getNextMatchedEvilCount(evil, matchedEvilCount, c, evilLPS);
      dp[i][matchedEvilCount][k1][k2] +=
          find(s1, s2, evil, i + 1, nextMatchedEvilCount, isS1Prefix && c == s1.charAt(i),
               isS2Prefix && c == s2.charAt(i), evilLPS);
      dp[i][matchedEvilCount][k1][k2] %= kMod;
    }
    return dp[i][matchedEvilCount][k1][k2];
  }

  // Get Longest Prefix also Suffix
  private int[] getLPS(final String s) {
    int[] lps = new int[s.length()];
    for (int i = 1, j = 0; i < s.length(); ++i) {
      while (j > 0 && s.charAt(j) != s.charAt(i))
        j = lps[j - 1];
      if (s.charAt(i) == s.charAt(j))
        lps[i] = ++j;
    }
    return lps;
  }

  // J := the next index we're trying to match with `currChar`
  private int getNextMatchedEvilCount(final String evil, int j, char currChar, int[] lps) {
    if (nextMatchedCount[j][currChar - 'a'] != null)
      return nextMatchedCount[j][currChar - 'a'];
    while (j > 0 && evil.charAt(j) != currChar)
      j = lps[j - 1];
    return nextMatchedCount[j][currChar - 'a'] = (evil.charAt(j) == currChar ? j + 1 : j);
  }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
class Solution:
  def findGoodStrings(self, n: int, s1: str, s2: str, evil: str) -> int:
    kMod = 1_000_000_007
    evilLPS = self._getLPS(evil)

    # NextMatchedCount(i, j) := # Next matched evil count given that there're
    # Already i matches with `evil` and the current is ('a' + j)
    # J := the next index we're trying to match with `currChar`
    @functools.lru_cache(None)
    def getNextMatchedEvilCount(j: int, currChar: str) -> int:
      while j > 0 and evil[j] != currChar:
        j = evilLPS[j - 1]
      return j + 1 if evil[j] == currChar else j

    # dp(i, j, k1, k2) := # Of good strings for s[i:] and there're already j
    # Matches with `evil`, where k1 is the 0//1 tight constraint for s1 and k2
    # Is the 0//1 tight constraint for s2
    @functools.lru_cache(None)
    def dp(i: int, matchedEvilCount: int, isS1Prefix: bool, isS2Prefix: bool) -> int:
      # s[:i] contains `evil`, so don't consider any ongoing strings
      if matchedEvilCount == len(evil):
        return 0
      # Run out of string, contributes one
      if i == n:
        return 1
      ans = 0
      minCharIndex = ord(s1[i]) if isS1Prefix else ord('a')
      maxCharIndex = ord(s2[i]) if isS2Prefix else ord('z')
      for charIndex in range(minCharIndex, maxCharIndex + 1):
        c = chr(charIndex)
        nextMatchedEvilCount = getNextMatchedEvilCount(matchedEvilCount, c)
        ans += dp(i + 1, nextMatchedEvilCount,
                  isS1Prefix and c == s1[i],
                  isS2Prefix and c == s2[i])
        ans %= kMod
      return ans

    return dp(0, 0, True, True)

  # Get Longest Prefix also Suffix
  def _getLPS(self, s: str) -> List[int]:
    lps = [0] * len(s)
    j = 0
    for i in range(1, len(s)):
      while j > 0 and s[j] != s[i]:
        j = lps[j - 1]
      if s[i] == s[j]:
        lps[i] = j + 1
        j += 1
    return lps