LeetCode 1286 Iterator for Combination (Python)

Posted by 小明MaxMing on August 13, 2020

题目

Design an Iterator class, which has:

  • A constructor that takes a string characters of sorted distinct lowercase English letters and a number combinationLength as arguments.
  • A function next() that returns the next combination of length combinationLength in lexicographical order.
  • A function hasNext() that returns True if and only if there exists a next combination.

解题思路

  1. 使用递归求出所有的结果,python可以使用yield from来节省空间
  2. 对于每个结果都对于一个二进制的数,如果这个数的1的个数为combinationLength,则是一个结果,从2^L-1开始遍历

代码

from math import gamma
class CombinationIterator:

    def __init__(self, characters: str, combinationLength: int):
        l = len(characters)
        self.count = int(gamma(l + 1) / gamma(l - combinationLength + 1) / gamma(combinationLength + 1))
        self.wg = self.wordgen(characters, combinationLength)

    def next(self) -> str:
        self.count -= 1
        return next(self.wg)

    def hasNext(self) -> bool:
        return self.count > 0
    
    def wordgen(self, characters, combinationLength):
        if combinationLength == 1:
            yield from characters
        else:
            for i in range(len(characters) - combinationLength + 1):
                yield from (characters[i] + tail for tail in self.wordgen(characters[i+1:], combinationLength - 1))
class CombinationIterator:

    def __init__(self, characters: str, combinationLength: int):
        self.s = characters[::-1]
        self.cur = (1 << len(characters)) - 1
        self.l = combinationLength

    def count(self, n):
        ret = 0
        while n:
            ret += 1
            n = n & (n - 1)
        return ret
    
    def next(self) -> str:
        while self.cur and self.count(self.cur) != self.l:
            self.cur -= 1
            
        res = []
        for i in range(len(self.s)):
            if (self.cur & (1<<i)) >> i:
                res.append(self.s[i])
        self.cur -= 1
        return ''.join(res)[::-1]

    def hasNext(self) -> bool:
        while self.cur and self.count(self.cur) != self.l:
            self.cur -= 1
        return self.cur > 0

视频讲解 YouTube<--欢迎点击订阅

视频讲解 bilibili<--欢迎点击订阅