Python, generators on a heap


  • 39

    Solution 1 ... ~550 ms (updated July 2016, originally was ~1570 ms)

    Using generators and heapq.merge. Too bad there's no itertools.unique.

    def nthSuperUglyNumber(self, n, primes):
        uglies = [1]
        def gen(prime):
            for ugly in uglies:
                yield ugly * prime
        merged = heapq.merge(*map(gen, primes))
        while len(uglies) < n:
            ugly = next(merged)
            if ugly != uglies[-1]:
                uglies.append(ugly)
        return uglies[-1]
    

    Solution 2 ... ~500 ms (updated July 2016, originally was ~1400 ms)

    Same thing done differently and it's a bit faster.

    def nthSuperUglyNumber(self, n, primes):
        uglies = [1]
        merged = heapq.merge(*map(lambda p: (u*p for u in uglies), primes))
        uniqed = (u for u, _ in itertools.groupby(merged))
        map(uglies.append, itertools.islice(uniqed, n-1))
        return uglies[-1]

  • 0

    short and clear!


  • 0
    S

    Vote for python!!!


  • 1
    S

    btw, may i ask why next(merged) function will not stop iteration in solution 1. I thought it only contains prime numbers... Please tell me why calling next will add new ugly to the merged function? Thank you so much!


  • 0

    Not sure how to answer that. But merged is not a function. It's an iterator. It merges the generators I give it. And each generator gives me the ugly numbers multiplied with that generator's prime.

    Are you maybe not familiar with generators? Here's a simpler one:

    >>> import itertools
    >>> cubes = (i**3 for i in itertools.count())
    >>> next(cubes)
    0
    >>> next(cubes)
    1
    >>> next(cubes)
    8
    >>> next(cubes)
    27
    >>> list(itertools.islice(cubes, 10))
    [64, 125, 216, 343, 512, 729, 1000, 1331, 1728, 2197]
    >>> for cube in cubes:
            print cube
    2744
    3375
    4096
    4913
    5832
    6859
    8000
    9261
    10648
    ...
    (continues until you stop it)

  • 0

    Thanks for sharing, is the inside function a closure?


  • 0

    @JayWong I'm not 100% sure about the terminology, but yes, I think it's a closure, capturing the uglies variable.


  • 0

    Sorry, why did you put a star * before map? You used it handle the StopIteration exception? I don't really understand :)


  • 1

    Just unpacking the list of generators returned by map, because heapq.merge needs them as separate parameters.


  • 2

    See the difference here:

    >>> import heapq
    >>> lists = [1, 3, 5], [2, 4, 6]
    
    >>> list(heapq.merge(lists))
    [[1, 3, 5], [2, 4, 6]]
    
    >>> list(heapq.merge(*lists))
    [1, 2, 3, 4, 5, 6]

  • 0

    I see! Thank you so much! The iterator in Python is pretty cool, can I understand it as dynamic? As long as we do not exhaust one iterator, when the dependent variables for its generator changed, the iterator would auto-update?


  • 0

    Sorry, I don't really understand what you're asking. Rephrase it maybe? Which iterator are you talking about?


  • 0

    merged in solution1. When you assign heapq.merge(*map(gen, primes)) to it, there is only 2,3,4 in the iterables. Why it can keeps generating without further updating?


  • 1

    Ah, ok. Well, the generators keep going as long as there is more in uglies, and uglies keeps getting extended.


  • 0

    Sorry for bothering you in this post. I think map(gen, primes) would return a list of iterators. Then heapq.merge(*map(gen, primes)) would give us an iterator of some iterable whose elements are iterators. What I think wrong?


  • 1

    Sounds like you're misunderstanding what heapq.merge does. It merges the given iterators of ints, giving you an iterator of ints, not of iterators.


  • 0
    A
    This post is deleted!

  • 0
    A

    Hi, Stefan, I am so confused with the generator. I add some debug sentences into your code.

    class Solution(object):
        def nthSuperUglyNumber(self, n, primes):
            """
            :type n: int
            :type primes: List[int]
            :rtype: int
            """
            uglies = [1]
            def gen(prime):
                print 'b'
                for ugly in uglies:
                    print 'c'
                    yield ugly * prime
            merged = heapq.merge(*map(gen, primes))
            while len(uglies) < n:
                print 'a'
                ugly = next(merged)
                print ugly
                if ugly != uglies[-1]:
                    uglies.append(ugly)
            return uglies[-1]
    

    The result for the test case 6, [2,3,5] is
    a
    b
    c
    b
    c
    b
    c
    2
    a
    c
    3
    a
    c
    4
    a
    c
    5
    a
    c
    6

    'b' was not printed in the last. What happened there and why 'c' can be printed without printing 'b'? And 'b' was printed three times in the beginning, but why it was not printed after that? I think merged is a generator, but how about the result for gen() function? Are they generators too? And What happened if different generators nested?


  • 0
    This post is deleted!

  • 1

    @annabel My gen is a generator function and returns a generator iterator. Here's a demo:

    >>> def count():
            i = 1
            while True:
                yield i
                i += 1
    
    >>> type(count)
    <type 'function'>
    >>> type(count())
    <type 'generator'>
    

    Now let's use it:

    >>> c = count()
    >>> next(c)
    1
    >>> next(c)
    2
    >>> next(c)
    3
    

    Note that I call the count function only once! After that, I'm calling the iterator it gave me. And that iterator does what's written in count, only initializes once and then does the loop, yielding the next value whenever I ask for one. That's why you only see one 'b' for each generator iterator.

    Check out generator and generator iterator.


Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.