Given an array of numbers, (or in this case a pandas series) return the longest wiggle subarray. A wiggle array is one such that the numbers alternate in direction- they go up and down strictly. I.e.
Input: [10, 22, 9, 33, 49, 50, 31, 60, 15]
Output: [49, 50, 31, 60, 15]
The array must be formed in sequence, i.e you cant form the array by skipping values.
My naive brute force approach can be found below- I'm having a tough time getting it to work. Any tips?
def zigzag(arr):
result = pd.Series([])
for i in range(len(arr)-2):
if arr[i:i+3] != sorted(arr[i:i+3]) and arr[i:i+3] != sorted(arr[i:i+3], reverse = True): # no three points increase or decrease
result = result.append(pd.Series(arr[i:i+3], index = list(range(i, i+3))))
result = result.loc[~result.index.duplicated(keep='first')]
else:
result = pd.Series([])
return result
You can implement a greedy solution which has a linear complexity by iterating and checking the zigzag condition, taking elements greedily until you reach an element which breaks the condition then you maximize the answer (the range you already have which doesn't break the condition) and repeat the operation starting from the current element (the one that broke the condition).
This works because if an element breaks the condition, then there is no better answer including the range before that element and the one after it.
I'm not on my machine to run the code but here's an example (follow the comments):
def zigzag(arr):
mx, cur, lst, op = 0, 0, -1, 0
st, en = 0, 0
ans = 0
for i, x in enumerate(arr):
# if its the first element we take it greedily
if lst == -1:
lst = x
st = i
# else if lst is the first element taken then decide which direction we need next
elif op == 0:
if x > lst:
op = 1
elif x < lst:
op = -1
# if current element is equal to last taken element we stop and start counting from this element
else:
# check if current range is better than previously taken one
if i-st > mx:
mx = i-st
ans = st
en = i
st = i
op = 0
lst = x
else:
# if direction meets the inequality take the current element
if (op == 1 and x < lst) or (op == -1 and x > lst):
op *= -1
lst = x
# otherwise, check if the current range is better than the previously taken one
else:
if i-st > mx:
mx = i-st
ans = st
en = i
st = i
lst = x
op = 0
# if starting index is greater than or equal to the last, then we have reached the end without checking for maximum range
if st >= en:
if len(arr)-st > en-ans:
ans = st
en = len(arr)
result = [arr[i] for i in range(ans, en)]
return result
print(zigzag([10, 22, 9, 33, 49, 50, 31, 60, 15]))
# [49, 50, 31, 60, 15]