본문 바로가기

코딩테스트

Daily Leetcoding Challenge 338. Counting Bits

오늘도 올려보는 리트코드 챌린지. Machine Learning 쪽 포지션들은 아무래도 파이썬을 주로 사용하는데, 때문에 코테를 볼 때 사용 언어를 파이썬으로 제한하는 경우가 있다. 그래서 앞으로는 파이썬으로도 코딩해보려고 한다. (난 타입스크립트를 더 자주 쓰고 싶은데..)

https://leetcode.com/problems/counting-bits

 

Counting Bits - LeetCode

Can you solve this real interview question? Counting Bits - Given an integer n, return an array ans of length n + 1 such that for each i (0 <= i <= n), ans[i] is the number of 1's in the binary representation of i.   Example 1: Input: n = 2 Output: [0,1,1

leetcode.com

직관

인자로 n이 주어졌을 때, 0부터 n까지 총 n+1 개 숫자의 이진법 표현에서 1의 개수를 리턴하는 문제다. i (0 <= i <= n)를 2로 나눈 나머지 또는 비트연산 i & 1 값이 끝자리(least significant digit)가 1 또는 0인지 알 수 있는 방법이다. 그다음 i를 2로 나누면(i / 2 또는 비트 시프트 i >> 1), 다음 자리의 비트가 끝자리로 이동한다.

 

접근법

각 숫자 i에 대해서 i가 0이 될 때까지 끝자리의 비트 값을 찾아 더하고, i를 i / 2 또는 i >> 1 업데이트한다.

 

복잡도

한 숫자에서 비트를 카운트하는 데 필요한 연산은 숫자를 0이 될 때까지 2로 나눠서 비트값을 찾기 때문에 logn이고, n+1개의 숫자가 있으므로 시간 복잡도는 O((n+1) * logn) = O(n * logn)이 된다. 공간 복잡도는 정답으로 리턴하는 배열을 제외하면 O(1)이다.

 

시간 복잡도가 O(n)인 솔루션도 존재한다. 비트 연산이 나누기 2라는 점을 활용하면 된다.

n = 5일 때 0부터 5까지의 2진법 표현을 보자:

0 --> 0

1 --> 1

2 --> 10

3 --> 11

4 --> 100

5 --> 101

4의 비트 카운트는 '10'(2) + 4 & 1과 같다. 즉 2의 비트 카운트 + 마지막 자리의 비트가 된다. 이것을 일반화하면 숫자 i의 비트 수는 i >> 1의 비트 수 + i & 1이다. 이 식을 활용하면 시간 복잡도 O(n)인 솔루션을 구할 수 있다.

 

코드

# O(nlogn)
class Solution:
    def countBits(self, n: int) -> List[int]:
        def countOne(n: int) -> int:
            ret = 0
            while n > 0:
                if n % 2 == 1:
                    ret += 1
                n = n // 2
            return ret

        result: List[int] = []
        for i in range (n+1):
            result.append(countOne(i))
        return result

# O(n)
class Solution:
    def countBits(self, n: int) -> List[int]:
        ans = [0]
        for i in range(1,n+1):
            ans.append(ans[i>>1] + (i &1))
        return ans
class Solution {
    public int[] countBits(int n) {
        int[] ret = new int[n+1];
        ret[0] = 0;
        for(int i=1; i<=n; i++) {
            int lastbit = i & 1;
            ret[i] = ret[i>>1] + lastbit;
        }
        return ret;
    }
}