TIL - 펜윅트리 문제편

LJB·2022년 7월 25일
0

TIL

목록 보기
3/3

지난 포스트에서 BIT에 대해서 알아보았다.

이제 이를 활용해서 리트코드 문제를 풀어보자

Count of Smaller Numbers After Self - LeetCode

문제 이해

문제는 간단하다

nums 라는 숫자가 들어있는 리스트가 주어지는데

각 숫자를 기준으로 그 숫자보다 오른쪽에 있는 값 중 그 숫자보다 작은 숫자의 개수를 값으로 가지는 리스트를 반환 하는 것이다.(이제 반환 할 리스트를 result라고 하겠다)

예를들어 nums가 [5,2,1,3] 이면

[3,1,0,0]을 반환 하는 것이다

  • 5 오른쪽에 있는 값 중 3개가 5보다 작다 (2,1,3)
  • 2 오른쪽에 있는 값 중 1개가 2보다 작다 (1)
  • 1,3오른쪽에 있는 값 중 1,3 각각의 숫자보다 작은 값은 없다

문제 풀이

이제 문제를 풀어야 할 시간이다

이 문제를 아주 단순히 생각하면

그냥 각 숫자마다 그 오른쪽에 있는 숫자를 다 돌면서 작은 값이 있는지 없는지 확인하고 총 개수를 result에 append해주면 될 것이다.

하지만 이것은 매우 비효율적인 풀이 방법으로

우리는 BIT를 여기에 적용 해야 할 것이다

BIT를 여기에 어떻게 적용시킬 수 있을까?

이 문제를 거꾸로 생각해보면

nums = [5,2,1,3]에서 생각해보면

result[0]을 구하는데 필요한 정보는 nums[1:4]에 대한 정보

result[1]을 구하는데 필요한 정보는 nums[2:4]에 대한 정보

result[2]를 구하는데 필요한 정보는 nums[3:4]에 대한 정보 이다

즉, nums를 반대로 돌면서 각 값에 대한 result값을 구해주고(BIT에서 getSum 메소드), 각 값을 update해주면 된다

이것을 코드로 쓰면

class Solution:
    def countSmaller(self, nums: List[int]) -> List[int]:

        ranking = {num : rank+1 for rank,num in enumerate(sorted(nums))}
        length = len(nums)
        result = []
        bi_tree = [0] * (length+1)

        def get_last_digit(i):
            return i&-i
        
        def update(i):
            while i <= length:
                bi_tree[i] += 1 
                i += get_last_digit(i)
                
        def get_sum(i):
            ret = 0
            while i:
                ret += bi_tree[i]
                i -= get_last_digit(i)
            return ret
        
        for num in nums[::-1]:
            result.append(get_sum(ranking[num]-1))
            update(ranking[num])
        return result[::-1]

이다.

잘 만든 자료구조의 강력함을 느낀 문제였다

profile
Better and better

0개의 댓글