Wiggle subsequence from leetcode

Question

This is the Wiggle Subsequence problem from leetcode.com:

A sequence of numbers is called a wiggle sequence if the differences between successive numbers strictly alternate between positive and negative. The first difference (if one exists) may be either positive or negative. A sequence with fewer than two elements is trivially a wiggle sequence.

For example, [1,7,4,9,2,5] is a wiggle sequence because the differences (6,-3,5,-7,3) are alternately positive and negative. In contrast, [1,4,7,2,5] and [1,7,4,5,5] are not wiggle sequences, the first because its first two differences are positive and the second because its last difference is zero.

Given a sequence of integers, return the length of the longest subsequence that is a wiggle sequence. A subsequence is obtained by deleting some number of elements (eventually, also zero) from the original sequence, leaving the remaining elements in their original order.

The solution is inspired by an \$O(n\log n)\$ answer to How to determine the longest increasing subsequence using dynamic programming?

class Solution(object):
 def wiggleMaxLength(self, nums):
 """
 :type nums: List[int]
 :rtype: int
 """
 if not nums or len(nums) == 1:
 return len(nums)
 current_num, prev_num = nums[0], nums[0]
 count, trend = 0, None
 for current_num in nums:
 if not trend:
 if current_num > prev_num:
 trend = "down_to_up"
 count += 2
 elif current_num < prev_num:
 trend = "up_to_down"
 count += 2
 elif trend == "down_to_up":
 if current_num < prev_num:
 count += 1
 trend = "up_to_down"
 elif trend == "up_to_down":
 if current_num > prev_num:
 prev_num = current_num
 count += 1
 trend = "down_to_up"
 prev_num = current_num
 if trend == None:
 count = 1
 return count

• IMCoins: According to how to ask, your question is not correct. The following code doesn't work for obvious reasons, but here is an example to prove it : [1, 7, 4, 8, 9, 2, 5] returns 6, and the correct answer is 4.
• Toby Speight: What is the return value supposed to mean? It's called count, but what's it a count of? (I tried following the "Question" link, but it's a mostly-empty page, so it didn't help).
• IMCoins: @TobySpeight It is supposed to be the length of the longuest wiggle subsequence, in given sequence. But OP's code doesn't work, and should go on stackoverflow probably.
• Gareth Rees: @IMCoins: The subsequence doesn't have to be contiguous: the challenge page gives the example [1,17,5,10,13,15,10,5,16,8] which has wiggle subsequences of length 7 for example [1,17,10,13,10,16,8]. I edited the question to quote some more of the problem, in particular "A subsequence is obtained by deleting some number of elements from the original sequence".
• IMCoins: Thanks for the insight @GarethRees . I'm then going to review the code. :)

Answer 1

1. The special case:

   if not nums or len(nums) == 1:
   

   can be simplified to:

   if not nums:
   

   because the case len(nums) == 1 is already handled correctly by the main code.

2. Instead of starting with count = 0, setting count += 2 when a trend is found for the first time, and having an adjustment at the end if no trend was found, you could start with count = 1, set count += 1 when a trend is found for the first time, and remove the adjustment.

3. There is no need to initialise current_num = nums[0] because this immediately gets overwritten by the for current_num in ... loop.

4. There is an unnecessary assignment of prev_num = current_num in one of the branches.

5. Using the strings "down_to_up" and "up_to_down" to represent the current trend is risky: if you made a typo and wrote "up_to_dwon" somewhere, then the code would be broken but Python would not detect this. It would be better to use global constants, for example

   _DOWN_TO_UP = "down to up"
   _UP_TO_DOWN = "up to down"
   

   and then if you make a typo like _UP_TO_DWON you will get a NameError from Python.

6. Now the main logic of the function looks like this:

   if not trend:
    if current_num > prev_num:
    trend = _DOWN_TO_UP
    count += 1
    elif current_num < prev_num:
    trend = _UP_TO_DOWN
    count += 1
   elif trend == _DOWN_TO_UP:
    if current_num < prev_num:
    trend = _UP_TO_DOWN
    count += 1
   elif trend == _UP_TO_DOWN:
    if current_num > prev_num:
    trend = _DOWN_TO_UP
    count += 1
   prev_num = current_num
   

   There is a lot of repetition here. In particular, on each branch we check to see if the trend has changed direction, and if it has, then we update trend and increment count. So we could refactor it so that we work out the new trend first, and then compare it with the old trend, like this:

   # Work out current trend
   if current_num > prev_num:
    current_trend = _DOWN_TO_UP
   elif current_num < prev_num:
    current_trend = _UP_TO_DOWN
   else: # current_num == prev_num
    current_trend = 0
   # Compare with previous trend
   if current_trend != 0 and current_trend != trend:
    trend = current_trend
    count += 1
   prev_num = current_num
   

7. In the case current_num == prev_num we know that the trend does not change, and so there can't be a change to trend, so it would make sense to test this first, and so avoid the test current_trend != 0, like this:

   if current_num != prev_num:
    # Work out current trend
    if current_num > prev_num:
    current_trend = _DOWN_TO_UP
    else:
    current_trend = _UP_TO_DOWN
    # Compare with previous trend
    if current_trend != trend:
    trend = current_trend
    count += 1
    prev_num = current_num
   

8. And now we can make a further simplification by using True instead of _DOWN_TO_UP and False instead of _UP_TO_DOWN:

   if current_num != prev_num:
    current_trend = current_num > prev_num
    if current_trend != trend:
    trend = current_trend
    count += 1
    prev_num = current_num
   

9. After making this change, it should be clear that prev_trend would be a better name than trend, for parallelism with prev_num.

   def wiggleMaxLength(self, nums):
    "Return length of longest wiggle subsequence of nums."
    if not nums:
    return 0
    count = 1
    prev_num = nums[0]
    prev_trend = None
    for current_num in nums:
    if current_num != prev_num:
    current_trend = current_num > prev_num
    if current_trend != prev_trend:
    prev_trend = current_trend
    count += 1
    prev_num = current_num
    return count
   

   It would make sense to stop here, as the revised code is already a lot simpler than the original. But if you are comfortable with Python iterators then you might want to read on to see how the functions and recipes in the itertools module can be used to further simplify the code.

10. First, we can avoid the current_num != prev_num test by using itertools.groupby to group identical numbers together:

    from itertools import groupby
    def wiggleMaxLength(self, nums):
     "Return length of longest wiggle subsequence of nums."
     if not nums:
     return 0
     nums = (n for n, _ in groupby(nums))
     prev_num = next(nums)
     count = 1
     prev_trend = None
     for current_num in nums:
     current_trend = current_num > prev_num
     if current_trend != prev_trend:
     prev_trend = current_trend
     count += 1
     prev_num = current_num
     return count
    

11. Now that we have an iterator over groups of identical numbers, we can use the pairwise recipe from the itertools documentation to iterate over adjacent pairs of numbers, avoiding the need to assign to prev_num.

    from itertools import groupby, tee
    def pairwise(iterable):
     "s -> (s0,s1), (s1,s2), (s2, s3), ..."
     a, b = tee(iterable)
     next(b, None)
     return zip(a, b)
    def wiggleMaxLength(self, nums):
     "Return length of longest wiggle subsequence of nums."
     if not nums:
     return 0
     nums = (n for n, _ in groupby(nums))
     count = 1
     prev_trend = None
     for prev_num, current_num in pairwise(nums):
     current_trend = current_num > prev_num
     if current_trend != prev_trend:
     prev_trend = current_trend
     count += 1
     return count
    

12. Now we can construct an iterator over the trends:

     trends = (cur > prev for prev, cur in pairwise(nums))
    

    and apply itertools.groupby to this iterator, thus grouping identical trends together and avoiding the need for current_trend and prev_trend. In fact, all we have to do now is count the groups of identical trends.

    def wiggleMaxLength(self, nums):
     "Return length of longest wiggle subsequence of nums."
     if not nums:
     return 0
     nums = (n for n, _ in groupby(nums))
     trends = (cur > prev for prev, cur in pairwise(nums))
     return 1 + sum(1 for _ in groupby(trends))
    

13. This could be written as a single expression:

    def wiggleMaxLength(self, nums):
     "Return length of longest wiggle subsequence of nums."
     return bool(nums) + sum(1 for _ in groupby(
     a > b for a, b in pairwise(n for n, _ in groupby(nums))))
    

    (But only if you like that kind of code.)