Skip to content

446. Arithmetic Slices II - Subsequence 👍

  • 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
class Solution {
 public:
  int numberOfArithmeticSlices(vector<int>& nums) {
    const int n = nums.size();
    int ans = 0;
    // dp[i][j] := the number of subsequences end in nums[j] nums[i]
    vector<vector<int>> dp(n, vector<int>(n));
    unordered_map<long, vector<int>> numToIndices;

    for (int i = 0; i < n; ++i)
      numToIndices[nums[i]].push_back(i);

    for (int i = 0; i < n; ++i)
      for (int j = 0; j < i; ++j) {
        const long target = nums[j] * 2L - nums[i];
        if (const auto it = numToIndices.find(target);
            it != numToIndices.cend())
          for (const int k : it->second)
            if (k < j)
              dp[i][j] += (dp[j][k] + 1);
        ans += dp[i][j];
      }

    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
class Solution {
  public int numberOfArithmeticSlices(int[] nums) {
    final int n = nums.length;
    int ans = 0;
    // dp[i][j] := the number of subsequences end in nums[j] nums[i]
    int[][] dp = new int[n][n];
    Map<Long, List<Integer>> numToIndices = new HashMap<>();

    for (int i = 0; i < n; ++i) {
      numToIndices.putIfAbsent((long) nums[i], new ArrayList<>());
      numToIndices.get((long) nums[i]).add(i);
    }

    for (int i = 0; i < n; ++i)
      for (int j = 0; j < i; ++j) {
        final long target = nums[j] * 2L - nums[i];
        if (numToIndices.containsKey(target))
          for (final int k : numToIndices.get(target))
            if (k < j)
              dp[i][j] += (dp[j][k] + 1);
        ans += dp[i][j];
      }

    return ans;
  }
}