Leetcode 446. Arithmetic Slices II - Subsequence

문제

Given an integer array nums, 
return the number of all the arithmetic subsequences of nums.

A sequence of numbers is called arithmetic 
if it consists of at least three elements and 
if the difference between any two consecutive elements is the same.

For example, [1, 3, 5, 7, 9], [7, 7, 7, 7], and [3, -1, -5, -9] are arithmetic sequences.
For example, [1, 1, 2, 5, 7] is not an arithmetic sequence.
A subsequence of an array is a sequence that can be formed by removing some elements (possibly none) of the array.

For example, [2,5,10] is a subsequence of [1,2,1,2,4,1,5,10].
The test cases are generated so that the answer fits in 32-bit integer.

주어진 배열에 대해 등차수열을 이루는 모든 길이가 3개 이상인 부분수열의 갯수를 구하시오

예시

Input: nums = [2,4,6,8,10]
Output: 7
Explanation: 모든 가능한 등차수열의 부분수열은 아래와 같다.
[2,4,6]
[4,6,8]
[6,8,10]
[2,4,6,8]
[4,6,8,10]
[2,4,6,8,10]
[2,6,10]
총 7개.

제한

• $1 <= nums.length <= 1000$
• $-2^{31} <= nums[i] <= 2^{31} - 1$

풀이

• 2차원 dp를 쓰되 row는 2개 이상의 등차수열을 이루는 마지막 숫자,
  col은 이루어진 등차수열의 차이를 key로, 등차수열의 갯수를 value로 가지는 Hashmap을 둔다.

예시로 [1, 2, 3, 4] 를 들어보겠다. 사용되는 hashmap 의 기본값은 반드시 0이다.

0번째 index로 끝나는 2개 이상의 수열은 만들어 질수 없기 때문에 1번째 index인 2부터 탐색을 시작한다.

1로 이룰수 있는 2개 이상의 등차수열은 0, 1 밖에 없기 때문에 dp에는 1:1 이 추가된다.

그 다음 3으로 끝나는 수열인데, 이를 찾기 위해 앞의 모든 index를 탐색한다.

앞에 올 index의 dp[index] 에 대해 dp[index][diff] + 1 만큼을 자기 자신 dp인 dp[self][diff] 에 추가한다.

근데 최종 결과는 길이가 3개 이상이여야 하는데, 추가된 1은 자기 자신을 포함한 길이 1이기에

결과는 dp[index][diff] 만큼만 더하면 된다.

class Solution:
    def numberOfArithmeticSlices(self, nums: List[int]) -> int:
        
        res = 0

        dp = [collections.defaultdict(int) for _ in nums]

        for i in range(1, len(nums)):
            for j in range(i):
                diff = nums[i] - nums[j]
                dp[i][diff] += 1 + dp[j][diff]
                res += dp[j][diff]

        return res