Python으로 최대 3번의 원소 변경 후 배열의 최댓값과 최솟값 최소 차이 구하기

문제 개요

nums라는 배열이 하나 주어져 있다고 가정해 봅시다. 한 번의 이동(move)마다 배열의 원소 하나를 임의의 값으로 변경할 수 있으며, 최대 3번의 이동을 수행한 뒤 배열에서 가장 큰 값(최댓값)과 가장 작은 값(최솟값)의 차이를 최소화해야 합니다.

예를 들어 입력이 nums = [3,7,2,12,16]이라면 결과는 1입니다. 배열을 [1,1,0,1,1] 형태로 바꾸면 최댓값은 1, 최솟값은 0이 되어 차이가 1이 되기 때문입니다.

접근 방법

배열의 길이가 4 이하라면 모든 원소를 같은 값으로 맞출 수 있으므로 차이는 항상 0입니다. 배열이 더 길다면, 먼저 배열을 정렬한 뒤 어느 쪽 끝에서 몇 개의 원소를 변경할지에 따른 네 가지 경우를 모두 검토합니다. 3번의 이동으로는 최대 3개의 원소만 바꿀 수 있으므로, 작은 쪽에서 i개, 큰 쪽에서 (3−i)개를 변경하는 경우를 생각하면 됩니다.

• nums의 크기가 4 이하이면 0을 반환합니다.

• nums를 오름차순으로 정렬합니다.

• ans를 무한대(infinity)로 초기화합니다.

• i를 0부터 3까지 반복하며 다음을 수행합니다.

  • mi := nums[i] — 변경하지 않고 남는 가장 작은 값

  • ma := nums[len(nums) − (3 − i + 1)] — 변경하지 않고 남는 가장 큰 값

  • ans := min(ma − mi, ans)

• ans를 반환합니다.

네 가지 경우를 정리하면 다음과 같습니다.

즉, i가 커질수록 왼쪽(작은 값)에서 더 많은 원소를 버리고 오른쪽(큰 값)에서는 더 적은 원소를 버립니다. 이 네 가지 경우 중 남는 범위가 가장 좁은 것이 곧 정답이 됩니다.

구현 예시

def solve(nums):
    if len(nums) <= 4:
        return 0
    nums.sort()
    ans = float("inf")
    for i in range(4):
        mi = nums[i]
        ma = nums[-(3-i+1)]
        ans = min(ma-mi, ans)
    return ans

nums = [3,7,2,12,16]
print(solve(nums))

입력

[3,7,2,12,16]

출력

1

복잡도 분석

정렬에 O(n log n)의 시간이 소요되고, 이후에는 네 번의 단순 비교만 수행하므로 전체 시간 복잡도는 O(n log n)입니다. 제자리(in-place) 정렬을 사용하면 추가 공간 복잡도는 O(1)입니다.