unicornication
unicornication

Reputation: 627

Python iterator not working as anticipated

I have the following code, I'm trying to compare some values and return the highest ones:

def high(it):
    it = iter(it)
    returnlist = []
    try:
        while True:
            one = next(it)
            two = next(it) 
            three = next(it)
            if three <= two and one <= two:
                returnlist.append(two)
    except StopIteration:
        pass
    return returnlist

It's half working, but not properly:

>>high([0,1,-1,3,8,4,3,5,4,3,8]) #(this works)
[1, 8, 5] 
>>high([5,2,4,9,6,1,3,8,0,7]) #(this doesn't work, should return [9,8]
[8] 
>>high(int(is_prime(p)) for p in irange(1,20)) #(doesn't work, returns four 1's - should return 5)
[1, 1, 1, 1]

Upvotes: 0

Views: 947

Answers (2)

wim
wim

Reputation: 362458

@Shashank comment is correct in that you are assuming the iterators are independent, when they are not. You can patch up your function using tee:

from itertools import tee

def high(it):

    it1,it2,it3 = tee(iter(it), 3)
    next(it2, None)
    next(it3, None); next(it3, None)

    returnlist = []
    try:
        while True:
            one = next(it1)
            two = next(it2) 
            three = next(it3)
            if three <= two and one <= two:
                returnlist.append(two)
    except StopIteration:
        pass
    return returnlist

What I think is a more pythonic way of implementing the same idea:

from itertools import tee, izip

def threesome(iterable):
    "s -> (s0,s1,s2), (s1,s2,s3), (s2,s3,s4), ..."
    a, b, c = tee(iterable, 3)
    next(b, None)
    next(c, None); next(c, None)
    return izip(a, b, c)

def high(it):
    return [x2 for x1, x2, x3 in threesome(it) if x2 == max(x1, x2, x3)]

By the way, I think your expected output for the last case is incorrect. You should see zeroes in the output aswell, because this will happen anytime you have three composite numbers in a row (e.g. 8,9,10 will satisfy your conditions).

Upvotes: 1

Raniz
Raniz

Reputation: 11113

What your code is doing is dividing the "list" into chunks of 3 and only comparing the numbers within each chunk. This means that you're only comparing the first 3 numbers with each other, then the next 3 with each other. What you want to do is to use a sliding window so that every number is compared with the ones next to it. You can do this by keeping track of previous values:

def high(lst):
    returnlist = []
    one = None
    two = None
    for three in lst:
        # If one is None here we haven't
        # reached the first set of numbers yet
        if not one is None: 
            if three <= two and one <= two:
                returnlist.append(two)
        # Update the sliding window
        one = two
        two = three
    return returnlist

Upvotes: 1

Related Questions