338. Counting Bitsを解いた

338. Counting Bits

leetcode.com

非負整数numが与えられるので、0からnumまでのiについて、2進数に変換したiに含まれる1の数を数えていき、それらを配列にしたものを返すことが求められる。

Example 1:

Input: 2

Output: [0,1,1]

注意点としてはFollow upとして時間計算量をO(n)で解くことを、ベストソリューションとして求められている点である。

Brute Force

初回の提出は、以下のようなコードになった。これは上記のFollow upを無視した解答であり、計算量はO(n*sizeof(integer))となる。

class Solution:
    def countBits(self, num: int) -> List[int]:
        res = []
        for i in range(num+1):
            res.append(str(bin(i)).count("1"))
        return res

実行速度の結果は92 msとなり、これは全体の45.81 %の参加者より速い結果となった。

ボトルネックとなっているのはcount()を使用している部分であると推察した。文字列に特定の文字がいくつ含まれるかという操作は、線形探索O(n)であるので、iの数が大きくなるにつれ計算量は増加してしまう。これをどうにかしたい。

Optimize

最終的に以下のようなコードとなった。任意の数字までを2進数表記したものを紙面上に書き出して眺めていると、bitの桁数が増加することで1の数は1つ増加することに加え、その桁数で表現できる数値範囲は、1つ前の桁数の数値範囲に依存することがわかった。直前の結果を再利用することができるので、DP(Dynamic Programming)の適用が考えられる。

class Solution:
    def countBits(self, num: int) -> List[int]:
        res = [0, 1]
        if num == 0:
            return [0]
        if num == 1:
            return res
        count = 2
        lim = 2
        ind = 0
        
        while count != num + 1:
            if lim == ind:
                lim *= 2
                ind = 0
            res.append(res[ind] + 1) 
            count += 1
            ind += 1
        return res

実行速度の結果は80 msとなり、これは全体の84.49 %の参加者より速い結果となった。劇的な実行速度の変化は生じていないが、他との比較値を見てみると十分に改善が行われたのではないだろうか。