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

Python으로 제한된 조건의 배열에서 주어진 인덱스의 최댓값 찾기

세 개의 값 n, index, maxSum이 주어졌다고 가정해 봅시다. 우리는 아래 조건을 모두 만족하는 배열 nums에서 nums[index]가 가질 수 있는 최댓값을 구해야 합니다.

  • 배열 nums의 크기는 정확히 n입니다.
  • 배열의 모든 요소는 양수여야 합니다.
  • 인접한 두 요소의 차이는 1 이하입니다. 즉, 모든 i(0 ≤ i < n-1)에 대해 |nums[i] − nums[i+1]| ≤ 1을 만족합니다.
  • 배열 요소들의 총합은 maxSum을 초과하지 않습니다.
  • nums[index]의 값이 최대화되어야 합니다.

예시

입력이 n = 6, index = 3, maxSum = 8이라면 출력은 2가 됩니다. 예를 들어 [1, 1, 2, 2, 1, 1] 배열은 위의 모든 조건을 만족하면서 총합이 정확히 8이고, 이때 nums[3] = 2로 주어진 제약 조건 하에서 만들 수 있는 최댓값이기 때문입니다.

접근 방법: 이진 탐색 활용

이 문제는 이진 탐색(Binary Search)을 사용하면 효율적으로 해결할 수 있습니다. 핵심 아이디어는 다음과 같습니다.

  • nums[index]의 값을 임의의 후보 값 v라고 가정합니다.
  • 인접한 요소의 차이가 1 이하라는 조건 때문에, index 위치에서 멀어질수록 값을 1씩 줄이는 것이(최솟값은 1) 총합을 최소화하는 방법입니다.
  • 따라서 v가 정해지면 등차수열의 합 공식을 이용해 배열 전체의 최소 가능 총합을 빠르게 계산할 수 있습니다.
  • 이 최소 총합이 maxSum 이하라면 v는 달성 가능한 값이므로 더 큰 값을 탐색하고, 초과한다면 더 작은 값을 탐색합니다.

구체적인 풀이 단계는 다음과 같습니다.

  1. 탐색 범위를 설정합니다. left := maxSum을 n으로 나눈 몫, right := maxSum + 1
  2. ans := 0으로 초기화합니다.
  3. left < right인 동안 다음을 반복합니다.
    • mid := left + (right − left)를 2로 나눈 몫
    • ind_l : 왼쪽 구간(인덱스 0 ~ index−1)의 최소 합을 등차수열 공식으로 계산합니다.
    • ind_r : 오른쪽 구간(인덱스 index ~ n−1, index 위치 포함)의 최소 합을 계산합니다.
    • 만약 ind_l + ind_r ≤ maxSum이라면, 해당 값이 가능하므로 ans := mid로 갱신하고 left := mid + 1로 더 큰 값을 시도합니다.
    • 그렇지 않으면 right := mid로 범위를 줄여 더 작은 값을 시도합니다.
  4. 반복이 끝나면 ans를 반환합니다.

이 알고리즘의 시간 복잡도는 이진 탐색에 의해 O(log maxSum)이며, 추가 공간 없이 상수 공간(O(1))으로 해결할 수 있습니다.

구현 예시

아래 구현을 통해 더 잘 이해해 보겠습니다.

def solve(n, index, maxSum):
   left, right = maxSum//n, maxSum+1
   ans = 0
   while(left<right):
      mid = left + (right-left)//2
      ind_l = (mid-1+max(1,mid-index))*min(index,mid-1)//2 + abs(min(0,mid-index-1))
      ind_r = (mid+max(1,mid-(n-index-1)))*min(n-index, mid)//2+ abs(min(0,mid-(n-index-1)-1))

      if ind_l + ind_r <=maxSum:
         ans = mid
         left = mid+1
      else:
         right = mid
   return ans

n = 6
index = 3
maxSum = 8
print(solve(n, index, maxSum))

입력

6, 3, 8

출력

2