Solution 1 ... ~550 ms (updated July 2016, originally was ~1570 ms)
Using generators and
heapq.merge. Too bad there's no
def nthSuperUglyNumber(self, n, primes): uglies =  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 =  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]
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)
@JayWong I'm not 100% sure about the terminology, but yes, I think it's a closure, capturing the
Sorry, why did you put a star
map? You used it handle the StopIteration exception? I don't really understand :)
Just unpacking the list of generators returned by
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 keeps getting extended.
Sorry for bothering you in this post. I think
map(gen, primes) would return a list of
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.
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 =  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
'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?
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.
Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.