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

파이썬으로 삽입 정렬 시 필요한 총 시프트(이동) 횟수 구하기 – 펜윅 트리 활용

문제 개요

배열 하나가 주어지고, 여기에 삽입 정렬(Insertion Sort)을 적용해야 한다고 가정해 보겠습니다. 삽입 정렬은 각 원소를 이미 정렬된 앞부분과 비교하여 자신의 올바른 위치로 한 칸씩 시프트(이동)시키는 방식으로 동작합니다. 이때 우리가 구해야 하는 것은 배열 전체가 정렬될 때까지 발생하는 총 시프트 횟수입니다. 결과값은 정수이며, 배열이 이미 정렬되어 있는 경우에는 0을 반환합니다.

예를 들어 입력이 다음과 같다면,

[4, 5, 3, 1, 2]

각 원소가 제자리에 삽입되는 과정은 아래와 같습니다.

[4, 5, 3, 1, 2] → 4 삽입 : 0회 시프트
[4, 5, 3, 1, 2] → 5 삽입 : 0회 시프트
[3, 4, 5, 1, 2] → 3 삽입 : 2회 시프트 (4, 5 이동)
[1, 3, 4, 5, 2] → 1 삽입 : 3회 시프트 (3, 4, 5 이동)
[1, 2, 3, 4, 5] → 2 삽입 : 3회 시프트 (3, 4, 5 이동)

따라서 총 시프트 횟수는 0 + 0 + 2 + 3 + 3 = 8이 됩니다.

핵심 아이디어: 반전(Inversion) 개수 세기

삽입 정렬에서 각 원소가 이동하는 횟수의 합은 배열 내 반전(inversion)의 총 개수와 같습니다. 반전이란 앞에 있는 원소가 뒤의 원소보다 큰 순서쌍을 의미합니다. 단순히 삽입 정렬을 실행하면서 이동 횟수를 직접 세면 O(n²)의 시간이 걸리지만, 펜윅 트리(Binary Indexed Tree, BIT)를 이용하면 각 원소를 처리할 때마다 '지금까지 등장한 원소 중 현재 값 이하인 원소의 개수'를 빠르게 집계할 수 있어 전체 문제를 O(n log M) 시간에 해결할 수 있습니다.

알고리즘 단계

  1. n := 입력 배열의 길이로 설정합니다.
  2. temp_arr := 크기가 1000001인 리스트를 0으로 초기화합니다. (펜윅 트리로 사용)
  3. ans := 0으로 초기화합니다.
  4. 입력 배열의 각 원소에 대해 다음을 수행합니다.
    • 현재 값을 val에 저장하고, val이 0보다 큰 동안 ans에 temp_arr[val]을 더한 뒤 val에서 val AND -val을 빼며 누적합을 조회합니다. (현재 값 이하인 선행 원소 개수 계산)
    • val을 현재 값으로 되돌린 후, val이 1000000 이하인 동안 temp_arr[val]에 1을 더하고 val에 val AND -val을 더하며 트리를 갱신합니다.
  5. 모든 원소를 처리한 후, ans = n × (n−1) ÷ 2 − ans 로 최종 반전 개수를 계산합니다.
  6. ans를 반환합니다.

파이썬 구현 예제

def solve(input_arr):
    length = len(input_arr)
    temp_arr = [0] * 1000001
    ans = 0
    for item in input_arr:
        val = item
        while val > 0:
            ans += temp_arr[val]
            val -= val & -val
        val = item
        while val <= 1000000:
            temp_arr[val] += 1
            val += val & -val
    ans = length * (length - 1) // 2 - ans
    return ans

print(solve([4, 5, 3, 1, 2]))

입력

[4, 5, 3, 1, 2]

출력

8

시간·공간 복잡도

각 원소마다 두 개의 while 루프가 최대 log(M)번씩 반복되므로 시간 복잡도는 O(n log M)이며, 크기 M+1의 트리 배열이 필요하므로 공간 복잡도는 O(M)입니다. 여기서 M은 배열 원소의 최댓값(본 예제에서는 1,000,000)입니다. 이 접근법 덕분에 원소 수가 많은 배열에서도 삽입 정렬의 총 이동 횟수를 효율적으로 구할 수 있습니다.