315. Count of Smaller Numbers After Self

[zeikar]

Problem link

https://leetcode.com/problems/count-of-smaller-numbers-after-self/

Problem Summary

배열에서 자신보다 뒤에 작은 수가 몇 개 나오는지 출력하는 문제.

Solution

처음엔 응?했지만 계속 보니 범위 문제라는 것이 보인다.
레드블랙트리에서 lower_bound 하거나 펜윅 트리를 쓰면 될 것 같다. 단, 레드 블랙 트리에서 lower_bound는 C++만 되는 것 같으니 간단한 펜윅 트리를 구현해보자.

https://www.acmicpc.net/blog/view/21

위 링크에서 간략하게 핵심을 잘 정리해 놓았다.

범위가 -10^4부터 10^4니까 20001 크기의 펜윅 트리가 필요하고, 값을 넣을 때 10001을 더해서 넣어주면 딱 맞는다. 자기 값을 펜윅 트리에 1을 더하게 업데이트 하고, 더 작은 수를 구할 때는 그냥 맨 앞(-10^4)부터 자신 -1 까지 더한 값이 자신 보다 작은 수의 합이 된다.

펜윅 트리의 자세한 설명은 블로그 링크를 참고 바람.

Source Code

from typing import List


class FenwickTree:
    def __init__(self, n: int):
        self.n = n
        self.tree = [0] * (n + 1)

    def lowbit(self, x: int) -> int:
        return x & -x

    def add(self, x: int, val: int) -> None:
        while x <= self.n:
            self.tree[x] += val
            x += self.lowbit(x)

    def sum(self, x: int) -> int:
        res = 0
        while x > 0:
            res += self.tree[x]
            x -= self.lowbit(x)
        return res


class Solution:
    def countSmaller(self, nums: List[int]) -> List[int]:
        fenwick = FenwickTree(20002)
        res = []
        for i in range(len(nums) - 1, -1, -1):
            res.append(fenwick.sum(nums[i] - 1 + 10001))
            fenwick.add(nums[i] + 10001, 1)

        return res[::-1]