# Python, generators on a heap

• 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]``````

• short and clear!

• Vote for python!!!

• 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!

• 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)``````

• Thanks for sharing, is the inside function a closure?

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

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

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

• 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]``````

• 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?

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

• `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?

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

• 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?

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

• This post is deleted!

• 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?

• This post is deleted!

• @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.

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