파이썬으로 요소 하나를 삭제한 뒤 최댓값·최솟값을 모두 포함하는 부분 리스트 개수 구하기

숫자로 이루어진 리스트 nums가 주어지고, 이 리스트에서 최대 한 개의 요소를 삭제할 수 있다고 가정해 보겠습니다. 이때 구해야 할 것은, 삭제가 끝난 뒤의 리스트에서 최댓값과 최솟값을 동시에 포함하는 연속 부분 리스트(sublist)의 최대 개수입니다.

예를 들어 입력이 nums = [3, 2, 6, 2, 4, 10]이라면 출력은 8이 됩니다. 값 10을 제거하면 리스트는 [3, 2, 6, 2, 4]가 되는데, 이때 최댓값 6과 최솟값 2를 모두 담고 있는 부분 리스트는 다음과 같이 여덟 개입니다.

[2, 6]
[6, 2]
[2, 6, 2]
[3, 2, 6]
[6, 2, 4]
[2, 6, 2, 4]
[3, 2, 6, 2]
[3, 2, 6, 2, 4]

해결 접근 방법

핵심 아이디어는 각 오른쪽 끝 인덱스마다 "최솟값과 최댓값을 모두 포함하면서 해당 인덱스에서 끝나는 부분 리스트"가 몇 개인지 세는 것입니다. 부분 리스트의 시작점은 최솟값과 최댓값 중 더 왼쪽에 있는 위치까지 어디든 잡을 수 있으므로, 매 순간 min(min_pos, max_pos) + 1개씩 누적하면 전체 개수를 구할 수 있습니다.

check() 함수의 동작

리스트 lst를 인자로 받는 함수 check()를 정의합니다.
mn := lst의 최솟값, mx := lst의 최댓값으로 설정합니다.
min_pos := None, max_pos := None, ret := 0으로 초기화합니다.
인덱스 i와 값 num을 기준으로 lst를 순회하며 다음을 수행합니다.
- num == mn이면 min_pos := i로 갱신합니다.
- num == mx이면 max_pos := i로 갱신합니다.
- min_pos 또는 max_pos가 아직 None이면 다음 반복으로 건너뜁니다.
- 그 외에는 ret := ret + min(min_pos, max_pos) + 1을 수행합니다.
순회가 끝나면 ret을 반환합니다.

메인 로직

nums의 길이가 1 이하이면 그 길이를 그대로 반환합니다.
ret := check(nums)로 기본값을 구합니다.
rem_cand를 [nums의 최솟값, nums의 최댓값] 순서로 확인하며 다음을 수행합니다.
- rem_cand가 nums에 딱 한 번만 등장한다면, 그 값을 지웠을 때 최댓값·최솟값 구성 자체가 달라질 수 있으므로 제거를 시도해볼 가치가 있습니다.
- idx := nums에서 rem_cand의 인덱스
- ret := max(ret, check(nums[:idx] + nums[idx+1:])) — 즉 해당 요소를 실제로 제거한 리스트에 대해 check()를 다시 실행하고, 더 큰 값을 취합니다.
최종 ret을 반환합니다.

참고로 최댓값이나 최솟값이 여러 번 등장한다면 하나를 지워도 전체 최댓값·최솟값은 변하지 않으므로, 유일하게 존재하는 경우만 추가로 확인하면 됩니다. 이 풀이의 시간 복잡도는 check()를 상수 번 호출하므로 O(n)입니다.

예제 코드

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

class Solution:
   def solve(self, nums):
      if len(nums) <= 1:
         return len(nums)
      def check(lst):
         mn, mx = min(lst), max(lst)
         min_pos, max_pos = None, None
         ret = 0
         for i, num in enumerate(lst):
            if num == mn:
               min_pos = i
            if num == mx:
               max_pos = i
            if min_pos is None or max_pos is None:
               continue
            ret += min(min_pos, max_pos) + 1
         return ret
      ret = check(nums)
      for rem_cand in [min(nums), max(nums)]:
         if nums.count(rem_cand) == 1:
            idx = nums.index(rem_cand)
            ret = max(ret, check(nums[:idx] + nums[idx + 1 :]))
      return ret
ob = Solution()
nums = [3, 2, 6, 2, 4, 10]
print(ob.solve(nums))

입력

[3, 2, 6, 2, 4, 10]