Skip to content

3388. Count Beautiful Splits in an Array

  • Time: $O(n^2)$
  • Space: $O(n^2)$
 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
class Solution {
 public:
  int beautifulSplits(vector<int>& nums) {
    const int n = nums.size();
    int ans = 0;
    // z[start][i] := the z array of nums[i..n) with respect to nums[start..n)
    vector<vector<int>> z;

    for (int start = 0; start < n; ++start)
      z.push_back(zFunction(nums, start));

    // nums1 | nums2 | nums3 = nums[0..i] | nums[i + 1..j] | nums[j + 1..n - 1]
    for (int i = 0; i < n - 2; ++i)
      for (int j = i + 1; j < n - 1; ++j)
        if ((j - i >= i + 1 &&
             z[0][i + 1] >= i + 1) ||  // nums1 is a prefix of nums2.
            z[i + 1][j + 1] >= j - i)  // nums2 is a suffix of nums3.
          ++ans;

    return ans;
  }

 private:
  // Returns the z array, where z[i] is the length of the longest prefix of
  // nums[i..n) which is also a prefix of nums[start..n).
  vector<int> zFunction(const vector<int>& nums, int start) {
    const int n = nums.size();
    vector<int> z(n);
    int l = 0;
    int r = 0;
    for (int i = 1 + start; i < n; ++i) {
      if (i < r)
        z[i] = min(r - i, z[i - l + start]);
      while (i + z[i] < n && nums[z[i] + start] == nums[i + z[i]])
        ++z[i];
      if (i + z[i] > r) {
        l = i;
        r = i + z[i];
      }
    }
    return z;
  }
};

 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
class Solution {
  public int beautifulSplits(int[] nums) {
    final int n = nums.length;
    int ans = 0;
    // z[start][i] := the z array of nums[i..n) with respect to nums[start..n)
    int[][] z = new int[n][];

    for (int start = 0; start < n; ++start)
      z[start] = zFunction(nums, start);

    // nums1 | nums2 | nums3 = nums[0..i] | nums[i + 1..j] | nums[j + 1..n - 1]
    for (int i = 0; i < n - 2; ++i)
      for (int j = i + 1; j < n - 1; ++j)
        if ((j - i >= i + 1 && z[0][i + 1] >= i + 1) || // nums1 is a prefix of nums2.
            z[i + 1][j + 1] >= j - i)                   // nums2 is a suffix of nums3.
          ++ans;

    return ans;
  }

  // Returns the z array, where z[i] is the length of the longest prefix of
  // nums[i..n) which is also a prefix of nums[start..n).
  private int[] zFunction(int[] nums, int start) {
    final int n = nums.length;
    int[] z = new int[n];
    int l = 0;
    int r = 0;
    for (int i = 1 + start; i < n; ++i) {
      if (i < r)
        z[i] = Math.min(r - i, z[i - l + start]);
      while (i + z[i] < n && nums[z[i] + start] == nums[i + z[i]])
        ++z[i];
      if (i + z[i] > r) {
        l = i;
        r = i + z[i];
      }
    }
    return z;
  }
}

 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
class Solution:
  def beautifulSplits(self, nums: list[int]) -> int:
    n = len(nums)
    # z[start][i] := the z array of nums[i..n) with respect to nums[start..n)
    z = [self._zFunction(nums, start)
         for start in range(n)]
    # nums1 | nums2 | nums3 = nums[0..i] | nums[i + 1..j] | nums[j + 1..n - 1]
    return sum((j - i >= i + 1 and z[0][i + 1] >= i + 1)  # nums1 is a prefix of nums2
               or z[i + 1][j + 1] >= j - i  # nums2 is a suffix of nums3.
               for i in range(n - 2)
               for j in range(i + 1, n - 1))

  def _zFunction(self, nums: list[int], start: int) -> list[int]:
    """
    Returns the z array, where z[i] is the length of the longest prefix of
    nums[i..n) which is also a prefix of nums[start..n).
    """
    n = len(nums)
    z = [0] * n
    l = start
    r = start
    for i in range(1 + start, n):
      if i < r:
        z[i] = min(r - i, z[i - l + start])
      while i + z[i] < n and nums[z[i] + start] == nums[i + z[i]]:
        z[i] += 1
      if i + z[i] > r:
        l = i
        r = i + z[i]
    return z