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

Python으로 '역전된 역전(Inverted Inversion)' 사중조 개수 구하기

문제 개요

숫자로 이루어진 리스트 nums가 주어졌다고 가정해 봅시다. 우리는 인덱스 조건 a < b < c < d를 만족하면서, 동시에 nums[a] < nums[b]이고 nums[c] > nums[d]인 사중조(quadruplet)의 개수를 구해야 합니다.

여기서 배열 nums는 1부터 N까지의 정수로 이루어진 순열(permutation)입니다.

예시로 이해하기

입력이 nums = [3, 4, 7, 6, 5]라면 출력은 5가 됩니다.

주어진 입력에서 찾을 수 있는 '역전된 역전'은 다음과 같습니다.

  • (3, 4, 7, 6)
  • (3, 4, 6, 5)
  • (3, 4, 7, 5)
  • (3, 7, 6, 5)
  • (4, 7, 6, 5)

해결 전략

핵심 아이디어는 중간 두 인덱스 b와 c를 기준으로 문제를 분리하는 것입니다. 특정 b에 대해 'b보다 앞에 있으면서 nums[b]보다 작은 값들의 개수의 누적합'과, 특정 c에 대해 'c보다 뒤에 있으면서 nums[c]보다 작은 값의 개수'를 미리 계산해 두면, 두 값을 곱해 전체 사중조 개수를 효율적으로 구할 수 있습니다.

알고리즘 단계

  1. 모듈로 상수 m := 10^9 + 7을 설정합니다.
  2. nums의 크기가 4 미만이면 0을 반환합니다.
  3. n := len(nums)로 설정합니다.
  4. sorted_ds 리스트를 만들고 nums의 마지막 원소를 삽입한 뒤 정렬합니다.
  5. ds_smaller_than_c := [0] * n 으로 초기화합니다.
  6. c를 n-2부터 0까지 역순으로 반복합니다.
    • ds_smaller_than_c[c] := sorted_ds에서 nums[c] - 1이 정렬 순서를 유지한 채 삽입될 수 있는 가장 오른쪽 위치 (즉, nums[c]보다 작은 원소의 개수)
    • sorted_ds 끝에 nums[c]를 삽입하고 정렬 상태를 유지합니다.
  7. quadruplet_count := 0으로 초기화합니다.
  8. sorted_as 리스트를 만들고 nums의 첫 번째 원소를 삽입한 뒤 정렬합니다.
  9. as_smaller_than_b_sum := 0으로 초기화합니다.
  10. b를 1부터 n-3까지 반복합니다.
    • as_smaller_than_b_sum에 sorted_as에서 nums[b] - 1이 삽입될 가장 오른쪽 위치를 더합니다.
    • 모듈로 연산을 적용합니다.
    • sorted_as에 nums[b]를 삽입하고 정렬 상태를 유지합니다.
    • quadruplet_count에 as_smaller_than_b_sum × ds_smaller_than_c[b + 1]을 더합니다.
    • 모듈로 연산을 적용합니다.
  11. quadruplet_count를 반환합니다.

여기서 bisect.bisect_right(sorted_list, x)는 정렬된 리스트에서 x가 삽입될 수 있는 가장 오른쪽 위치를 반환하므로, 현재 값보다 작은 원소의 개수를 로그 시간 안에 효율적으로 구할 수 있습니다.

구현 예제

import bisect
MOD = 10 ** 9 + 7

class Solution:
   def solve(self, nums):
      if len(nums) < 4:
         return 0
      n = len(nums)
      sorted_ds = list([nums[-1]])
      sorted_ds.sort()
      ds_smaller_than_c = [0] * n
      for c in range(n - 2, -1, -1):
         ds_smaller_than_c[c] = bisect.bisect_right(sorted_ds, nums[c] - 1)
         sorted_ds.append(nums[c])
         sorted_ds.sort()
      quadruplet_count = 0
      sorted_as = list([nums[0]])
      sorted_as.sort()
      as_smaller_than_b_sum = 0
      for b in range(1, n - 2):
         as_smaller_than_b_sum += bisect.bisect_right(sorted_as, nums[b] - 1)
         sorted_as.sort()
         as_smaller_than_b_sum %= MOD
         sorted_as.append(nums[b])
         sorted_as.sort()
         quadruplet_count += as_smaller_than_b_sum * ds_smaller_than_c[b + 1]
         quadruplet_count %= MOD
      return quadruplet_count

ob = Solution()
print(ob.solve([3, 4, 7, 6, 5]))

실행 결과

입력:

[3, 4, 7, 6, 5]

출력:

5

복잡도 분석

각 반복마다 리스트 정렬과 이분 탐색이 수행되므로 시간 복잡도는 O(N² log N), 공간 복잡도는 O(N)입니다. 입력 크기가 매우 큰 경우에는 펜윅 트리(Fenwick Tree)나 세그먼트 트리를 활용하면 O(N log N)까지 성능을 개선할 수 있습니다.