Problem 446

Question statement:

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.

Solution

This one was unreasonably hard, and even utilizing hints and pointers in the discussion section, this took me a while to figure out.

Main idea:

This problem reeks of dynamic programming. The natural first step is to think about successively finding the number of arithmetic subsequences in nums[:i]. The second obeservation is that we also need to memoize the common difference of each arithmetic sequence ending at index i. This is necessary if at a later point we spot another element j with the same common difference with i, so we can increase the length of this subsequence by 1. Moreover, a valid subsequence must have length at least three. The last two observations enable two realizations:

1. The dp ‘array’ must be multidimensional, in fact, one dimension for each valid common difference.
2. The values stored should be lengths of subsequences ending at index dp[i][common_diff] not the number of subsequences.

To count the number of subsequences, we can just update a counter everytime we encounter a subsequence of length at least 3.

To programattically approach the above realizations:

We can let dp be an array of dictionaries (this is a really cool dp idea, that is not mine. I was inspired by this other solution that was posted). This way dp[i][diff] will store the lengths of all arithmetic subsequence with common difference diff ending at index i. Now for each i, we will scan every previous element of nums and calculate diff = nums[i]-nums[j]. Now, if dp[j][diff] is non-zero, then there is already an arithmetic subsequence of common difference diff ending at index j. In this case, the length of my new sequence is simply dp[i][diff] = 1 + dp[j][diff]. Now, I check if dp[i][diff] is at least 2 (indicating the existence of subsequence of length at least 3), if so then I need to update my counter. Here’s where the next twist comes: counter needs to be updated by by dp[j][diff] (as opposed to 1), because after appending an arithmetic subsequence ending at j to nums[i], we now have dp[j][diff] many new subsequences that end at index i (count backwards from i adding one element of the subsequence at a time).

For instance: Let nums = [9, 7, 2, 4, 6, 8, 10], i = 6 and j = 5. In this case diff = 2. So, dp[6][2] = 1 + dp[5][2]. Here dp[5][2] would have memoized to be 3 (len([2,4,6,8])-1). Now to update the counter, we observe that in this example dp[6][2] is capturing only the arithmetic subsequence [2,4,6,8, 10]. This subsequence comes with it, even smaller arithmetic subsequences with the same difference. Namely, [6,8,10], [4,6,8,10] and [2,4,6,8,10] which is precisely dp[5][2].

Note:

• The subsequence [2,4,6] for instance, would have been accounted for in the case when i=4, so we only need to care about subsequences ending in 10.
• If there is another 8 elsewhere at index k, then diff = 2 in this case too, and during the iteration j=k, we will add the corresponding length to dp[6][2].

Code:

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

Comments:

Time complexity $\texttt{O(n}^2)$ where n is the length of nums

The above code ran in 632ms, beats 58.22% of users. It uses 72MB of memory, beating 27.11% of users.