Python으로 정렬된 배열에서 절대 차이의 합 효율적으로 구하기

문제 개요

비내림차순(오름차순)으로 정렬된 배열 nums가 주어졌다고 가정해 봅시다. 이때 nums와 길이가 같은 배열 result를 만들어야 하며, result[i]에는 nums[i]와 배열 안의 다른 모든 원소들 사이의 절대 차이의 합이 저장되어야 합니다.

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

• |5−5| + |5−7| + |5−12| = 0 + 2 + 7 = 9
• |7−5| + |7−7| + |7−12| = 2 + 0 + 5 = 7
• |12−5| + |12−7| + |12−12| = 7 + 5 + 0 = 12

접근 방법

배열이 이미 정렬되어 있다는 점을 활용하면, 모든 쌍을 일일이 비교하지 않고도 O(n) 시간 안에 답을 구할 수 있습니다. 핵심 아이디어는 인덱스가 하나씩 이동할 때마다 이전 누적값을 바탕으로 결과를 갱신하는 것입니다.

풀이 과정은 다음과 같습니다.

• res : 결과를 담을 새로운 리스트
• s : 누적 절대 차이 합, 초기값 0
• n : nums의 크기
• i를 1부터 n−1까지 반복하며 s += nums[i] − nums[0] → 첫 번째 원소(인덱스 0)에 대한 절대 차이 합을 구합니다.
• s를 res 끝에 추가합니다.
• 다시 i를 1부터 n−1까지 반복하며 다음을 수행합니다.
  • diff = nums[i] − nums[i−1]
  • s += diff * i → 현재 위치보다 왼쪽에 있는 i개의 원소들은 diff만큼 멀어지므로 더해줍니다.
  • s -= diff * (n − i) → 오른쪽에 있는 (n−i)개의 원소들은 diff만큼 가까워지므로 빼줍니다.
• 갱신된 s를 res에 추가한 뒤, 최종적으로 res를 반환합니다.

파이썬 구현 예제

아래 구현을 통해 동작 방식을 더 쉽게 이해할 수 있습니다.

def solve(nums):
    res = []
    s = 0
    n = len(nums)
    for i in range(1, n):
        s += nums[i] - nums[0]
    res.append(s)
    for i in range(1, n):
        diff = nums[i] - nums[i-1]
        s += diff * i
        s -= diff * (n - i)
        res.append(s)
    return res

nums = [5, 7, 12]
print(solve(nums))

입력

[5, 7, 12]

출력

[9, 7, 12]

시간 복잡도

배열을 두 번 선형 순회할 뿐이므로 전체 시간 복잡도는 O(n)이며, 결과 배열을 제외한 추가 공간은 O(1)입니다. 반면 모든 원소 쌍을 직접 비교하는 브루트포스 방식은 O(n²)이 소요되므로, 정렬된 배열에서는 위의 누적합 기법이 훨씬 효율적입니다.