Skip to content

3261. Count Substrings That Satisfy K-Constraint II 👍

  • Time: $O(n)$
  • Space: $O(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
class Solution {
 public:
  vector<long long> countKConstraintSubstrings(string s, int k,
                                               vector<vector<int>>& queries) {
    const int n = s.size();
    vector<long long> ans;
    vector<int> count(2);
    // leftToRight[l] : = the maximum right index r s.t.s[l..r] is valid
    vector<int> leftToRight(n);
    // rightToLeft[r] := the minimum left index l s.t. s[l..r] is valid
    vector<int> rightToLeft(n);
    // prefix[i] := the number of valid substrings ending in [0..i - 1].
    vector<long> prefix{0};

    for (int l = 0, r = 0; r < n; ++r) {
      ++count[s[r] - '0'];
      while (count[0] > k && count[1] > k)
        --count[s[l++] - '0'];
      rightToLeft[r] = l;
    }

    count = vector<int>(2);

    for (int l = n - 1, r = n - 1; l >= 0; --l) {
      ++count[s[l] - '0'];
      while (count[0] > k && count[1] > k)
        --count[s[r--] - '0'];
      leftToRight[l] = r;
    }

    for (int r = 0; r < n; ++r)
      prefix.push_back(prefix.back() + r - rightToLeft[r] + 1);

    for (const vector<int>& query : queries) {
      const int l = query[0];
      const int r = query[1];
      long numValidSubstrings = 0;
      if (r > leftToRight[l]) {
        // If r is beyond leftToRight[l], compute the number of valid substrings
        // from l to leftToRight[l] and add the number of valid substrings
        // ending in [leftToRight[l] + 1..r].
        //
        // prefix[r + 1] := the number of valid substrings ending in [0..r].
        // prefix[leftToRight[l] + 1] := the number of valid substrings ending
        // in [0..leftToRight].
        // => prefix[r + 1] - prefix[leftToRight[l] + 1] := the number of valid
        // substrings ending in [leftToRight[l] + 1..r].
        const int sz = leftToRight[l] - l + 1;
        numValidSubstrings =
            (sz * (sz + 1)) / 2 + (prefix[r + 1] - prefix[leftToRight[l] + 1]);
      } else {
        // If r is within the range of leftToRight[l], compute the number of
        // valid substrings directly from l to r.
        const int sz = r - l + 1;
        numValidSubstrings = (sz * static_cast<long>(sz + 1)) / 2;
      }
      ans.push_back(numValidSubstrings);
    }

    return ans;
  }
};

 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
class Solution {
  public long[] countKConstraintSubstrings(String s, int k, int[][] queries) {
    final int n = s.length();
    long[] ans = new long[queries.length];
    int[] count = new int[2];
    // leftToRight[l] : = the maximum right index r s.t.s[l..r] is valid
    int[] leftToRight = new int[n];
    // rightToLeft[r] := the minimum left index l s.t. s[l..r] is valid
    int[] rightToLeft = new int[n];
    // prefix[i] := the number of valid substrings ending in [0..i - 1].
    long[] prefix = new long[n + 1];

    for (int l = 0, r = 0; r < n; ++r) {
      ++count[s.charAt(r) - '0'];
      while (count[0] > k && count[1] > k)
        --count[s.charAt(l++) - '0'];
      rightToLeft[r] = l;
    }

    Arrays.fill(count, 0);

    for (int l = n - 1, r = n - 1; l >= 0; --l) {
      ++count[s.charAt(l) - '0'];
      while (count[0] > k && count[1] > k)
        --count[s.charAt(r--) - '0'];
      leftToRight[l] = r;
    }

    for (int r = 0; r < n; ++r)
      prefix[r + 1] = prefix[r] + r - rightToLeft[r] + 1;

    for (int i = 0; i < queries.length; ++i) {
      final int l = queries[i][0];
      final int r = queries[i][1];
      long numValidSubstrings = 0;
      if (r > leftToRight[l]) {
        // If r is beyond leftToRight[l], compute the number of valid substrings
        // from l to leftToRight[l] and add the number of valid substrings
        // ending in [leftToRight[l] + 1..r].
        //
        // prefix[r + 1] := the number of valid substrings ending in [0..r].
        // prefix[leftToRight[l] + 1] := the number of valid substrings ending
        // in [0..leftToRight].
        // => prefix[r + 1] - prefix[leftToRight[l] + 1] := the number of valid
        // substrings ending in [leftToRight[l] + 1..r].
        final int sz = leftToRight[l] - l + 1;
        numValidSubstrings = (sz * (sz + 1)) / 2 + (prefix[r + 1] - prefix[leftToRight[l] + 1]);
      } else {
        final int sz = r - l + 1;
        numValidSubstrings = (sz * (long) (sz + 1)) / 2;
      }
      ans[i] = numValidSubstrings;
    }

    return ans;
  }
}

 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
class Solution:
  def countKConstraintSubstrings(
      self,
      s: str,
      k: int,
      queries: list[list[int]]
  ) -> list[int]:
    n = len(s)
    ans = []
    count = [0, 0]
    # leftToRight[l] := the maximum right index r s.t. s[l..r] is valid
    leftToRight = [0] * n
    # rightToLeft[r] := the minimum left index l s.t. s[l..r] is valid
    rightToLeft = [0] * n

    l = 0
    for r in range(n):
      count[int(s[r])] += 1
      while min(count) > k:
        count[int(s[l])] -= 1
        l += 1
      rightToLeft[r] = l

    count = [0, 0]
    r = n - 1
    for l in reversed(range(n)):
      count[int(s[l])] += 1
      while min(count) > k:
        count[int(s[r])] -= 1
        r -= 1
      leftToRight[l] = r

    # prefix[i] := the number of valid substrings ending in [0..i - 1].
    prefix = list(itertools.accumulate((r - l + 1
                                       for r, l in enumerate(rightToLeft)),
                                       initial=0))

    for l, r in queries:
      if r > leftToRight[l]:
        # If r is beyond leftToRight[l], compute the number of valid substrings
        # from l to leftToRight[l] and add the number of valid substrings
        # ending in [leftToRight[l] + 1..r].
        #
        # prefix[r + 1] := the number of valid substrings ending in [0..r].
        # prefix[leftToRight[l] + 1] := the number of valid substrings ending
        # in [0..leftToRight].
        # => prefix[r + 1] - prefix[leftToRight[l] + 1] := the number of valid
        # substrings ending in [leftToRight[l] + 1..r].
        sz = leftToRight[l] - l + 1
        numValidSubstrings = sz * (sz + 1) // 2 + (
            prefix[r + 1] - prefix[leftToRight[l] + 1])
      else:
        # If r is within the range of leftToRight[l], compute the number of
        # valid substrings directly from l to r.
        sz = r - l + 1
        numValidSubstrings = sz * (sz + 1) // 2
      ans.append(numValidSubstrings)

    return ans