[LeetCode/Python3] Arithmetic Slices II - Subsequence

Problem

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개 이상이어야 한다.

Solution

1차 시도

class Solution:

    def numberOfArithmeticSlices(self, nums: List[int]) -> int:
        nums.sort()
        n = len(nums)
        answer = 0

        def backtracking(prev, sub_sequence, max_n):
            nonlocal answer
            if len(sub_sequence) > 2:
                answer += 1
            if prev + 1 == max_n:
                return
            for curr in range(prev + 1, max_n):
                if len(sub_sequence) < 2 or sub_sequence[-1] - sub_sequence[-2] == nums[curr] - sub_sequence[-1]:
                    sub_sequence.append(nums[curr])
                    backtracking(curr, sub_sequence, max_n)
                    sub_sequence.pop(-1)
        

        for i in range(n):
            backtracking(i, [nums[i]], n)
                    
        return answer

• 시간 초과 발생

2차 시도 (Discussion 참조)

from collections import defaultdict
class Solution:

    def numberOfArithmeticSlices(self, nums: List[int]) -> int:
        n = len(nums)
        total_count = 0
        dp = [defaultdict(int) for i in range(n)]
        for i in range(1, n):
            for j in range(i):
                diff = nums[i] - nums[j]
                dp[i][diff] += 1
                if diff in dp[j]:
                    dp[i][diff] += dp[j][diff]
                    total_count += dp[j][diff]
        return total_count

• Dynamic Programming을 이용해 풀이
  - 1부터 각 nums[i] 별로 diff에 따른 arithmetic subsequences의 개수를 구한다.
  - i보다 큰 수에서 diff에 따른 arithmetic subsequences의 개수는 nums[i - diff][diff]의 개수를 더한 것과 같다.
• dp[j][diff]는 0부터 j의 범위 내에서 공차가 diff인 수열의 개수이다
• 각 (i, j) 쌍이 있을 때 diff = nums[i] - nums[j]가 된다.
• dp[i][diff]를 1 증가시키고 새로운 서브 등차수열을 구한다
  - 만약 dp[j][diff]가 존재하면 서브 등차수열이 존재하는 것이므로 dp[j][diff]의 값을 dp[i][diff]와 total_count에 더한다
• 탐색을 끝마치면 total_count가 arithmetic subsequencs의 개수를 담게 된다.