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

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.)