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

Python으로 리스트 균형 맞추기: 양쪽 끝에서 제거해야 하는 최소 삭제 횟수 구하기

0과 1로만 이루어진 리스트가 있다고 가정해 보겠습니다. 우리는 리스트의 앞쪽 또는 뒤쪽에서만 값을 삭제할 수 있으며, 삭제가 끝난 뒤 남은 리스트에 포함된 0과 1의 개수가 서로 같아지도록 만들어야 합니다. 이 문제의 목표는 이때 필요한 최소 삭제 횟수를 구하는 것입니다.

예를 들어 입력이 nums = [1, 1, 1, 0, 0, 1]이라면 정답은 2입니다. 맨 앞의 1 하나와 맨 뒤의 1 하나를 각각 삭제하면 [1, 1, 0, 0]이 되어 1 두 개와 0 두 개로 균형이 맞기 때문입니다.

문제 해결 아이디어

핵심은 "균형이 맞는 가장 긴 연속 구간"을 찾는 것입니다. 전체 리스트 길이에서 이 구간의 길이를 빼면, 양쪽 끝에서 제거해야 하는 최소 횟수를 바로 알 수 있습니다.

이를 위해 누적 합(prefix sum) 기법을 활용합니다. 0을 만나면 누적 합(currSum)을 1 감소시키고, 1을 만나면 1 증가시킵니다. 그러면 서로 다른 두 인덱스에서 currSum 값이 같다는 것은, 그 사이 구간에 0과 1의 개수가 동일하다는 의미가 됩니다.

알고리즘 단계

  • longest := 0으로 초기화 (균형이 맞는 가장 긴 구간의 길이)
  • d := 딕셔너리를 생성하고, 키 0에 값 -1을 미리 저장
  • currSum := 0으로 초기화
  • i를 0부터 nums의 길이 - 1까지 반복:
    • nums[i]가 0이면 currSum을 1 감소, 그렇지 않으면 1 증가
    • currSumd에 이미 존재하면, longestlongesti - d[currSum] 중 더 큰 값으로 갱신
    • 존재하지 않으면 d[currSum] = i로 저장
  • 마지막으로 len(nums) - longest를 반환

키 0에 -1을 미리 넣어두는 이유는, 리스트의 처음부터 특정 인덱스까지 이미 균형이 맞는 경우도 처리하기 위함입니다.

구현 예제

class Solution:
    def solve(self, nums):
        longest = 0
        d = {0: -1}
        currSum = 0
        for i in range(len(nums)):
            if nums[i] == 0:
                currSum -= 1
            else:
                currSum += 1
            if currSum in d:
                longest = max(longest, i - d[currSum])
            else:
                d[currSum] = i
        return len(nums) - longest

ob = Solution()
nums = [1, 1, 1, 0, 0, 1]
print(ob.solve(nums))

입력

[1, 1, 1, 0, 0, 1]

출력

2

복잡도 분석

시간 복잡도: O(n) — 리스트를 한 번만 순회하므로 매우 효율적입니다.
공간 복잡도: O(n) — 최악의 경우 딕셔너리에 n개의 키가 저장될 수 있습니다.