Computer >> 컴퓨터 >  >> 프로그래밍 >> Python

파이썬으로 리스트의 각 요소 오른쪽에 있는 더 작은 요소 개수 구하기

숫자로 이루어진 리스트 nums가 주어졌을 때, 새로운 리스트를 만들어야 합니다. 새 리스트의 각 요소는 원본 리스트에서 해당 요소의 오른쪽에 위치한 더 작은 요소의 개수를 나타냅니다.

예를 들어 입력이 nums = [4, 5, 9, 7, 2]라면 출력은 [1, 1, 2, 1, 0]이 됩니다. 그 이유는 다음과 같습니다.

  • 4의 오른쪽에는 더 작은 요소가 1개 있습니다 (2)
  • 5의 오른쪽에는 더 작은 요소가 1개 있습니다 (2)
  • 9의 오른쪽에는 더 작은 요소가 2개 있습니다 (7, 2)
  • 7의 오른쪽에는 더 작은 요소가 1개 있습니다 (2)
  • 2의 오른쪽에는 더 작은 요소가 없습니다

해결 접근 방법

이 문제는 리스트를 오른쪽에서 왼쪽으로 순회하면서, 지금까지 본 요소들을 정렬된 상태로 유지하는 보조 리스트를 활용하면 효율적으로 해결할 수 있습니다. 구체적인 단계는 다음과 같습니다.

  1. 결과를 저장할 리스트 res와, 이미 확인한 요소들을 정렬된 상태로 관리할 리스트 inc를 생성합니다.
  2. nums가 빌 때까지 다음을 반복합니다.
    • num := nums의 마지막 요소를 꺼냅니다.
    • inc에서 num이 삽입될 가장 왼쪽 인덱스를 구합니다. 이 값은 현재까지 확인한 요소 중 num보다 작은 요소의 개수와 같으므로, res의 끝에 추가합니다.
    • numinc에 정렬된 상태를 유지하며 삽입합니다.
  3. 오른쪽부터 처리했기 때문에 res를 뒤집어서 반환합니다.

예제 코드

파이썬의 bisect 모듈을 사용하면 정렬된 리스트에서의 탐색과 삽입을 O(log n) 및 효율적으로 처리할 수 있습니다.

import bisect
class Solution:
    def solve(self, nums):
        res, inc = [], []
        while nums:
            num = nums.pop()
            res.append(bisect.bisect_left(inc, num))
            bisect.insort(inc, num)
        return res[::-1]
ob = Solution()
nums = [4, 5, 9, 7, 2]
print(ob.solve(nums))

입력

[4, 5, 9, 7, 2]

출력

[1, 1, 2, 1, 0]

동작 원리 살펴보기

bisect.bisect_left(inc, num)은 정렬된 리스트 inc에서 num이 들어갈 위치를 찾아주는데, 이 위치의 인덱스 값이 곧 inc에 있는 요소 중 num보다 작은 요소의 개수입니다. 또한 bisect.insort()는 정렬 순서를 깨뜨리지 않으면서 요소를 삽입해 주므로, 매번 리스트를 다시 정렬할 필요가 없습니다.

이 알고리즘의 시간 복잡도는 요소 하나당 삽입에 최대 O(n)이 소요되므로 전체적으로 O(n²)입니다. 단순한 이중 반복문(O(n²))보다 상수 배가 작아 실질적으로 더 빠르며, 코드도 간결합니다. 만약 더 큰 입력에 대해서는 펜윅 트리(Fenwick Tree)나 세그먼트 트리를 활용해 O(n log n)으로 개선할 수 있습니다.