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

파이썬으로 리스트 양 끝에서 K개의 숫자를 제거해 최대 합 구하기

숫자로 이루어진 리스트 nums와 정수 k가 주어졌을 때, 리스트의 왼쪽 또는 오른쪽 끝에서 정확히 k번 원소를 제거한다고 가정해 봅시다. 이때 제거할 수 있는 원소들의 합 중 최댓값을 구하는 것이 목표입니다.

문제 예시

예를 들어 입력이 다음과 같다고 해보겠습니다.

  • nums = [2, 4, 5, 3, 1]
  • k = 2

이 경우 출력은 6이 됩니다. 왼쪽 끝에서 2와 4를 차례로 제거하면 합이 6이 되기 때문입니다.

풀이 접근 방법

이 문제는 슬라이딩 윈도우(Sliding Window) 기법으로 효율적으로 해결할 수 있습니다. 처음에는 왼쪽에서 k개를 모두 선택한 상태에서 시작한 뒤, 한 번에 하나씩 왼쪽 원소를 빼고 오른쪽 원소를 더하면서 가능한 모든 조합을 확인합니다.

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

  1. window := 인덱스 0부터 k-1까지 원소들의 합
  2. ans := window (현재까지의 최댓값)
  3. i를 1부터 k까지 반복하면서:
    • window에서 nums[k - i]를 뺌 (왼쪽 원소 제거)
    • window에 nums[-i]를 더함 (오른쪽 원소 추가)
    • ans := ans와 window 중 큰 값으로 갱신
  4. ans 반환

파이썬 구현 코드

class Solution:
    def solve(self, nums, k):
        window = sum(nums[:k])
        ans = window
        for i in range(1, k + 1):
            window -= nums[k - i]
            window += nums[-i]
            ans = max(ans, window)
        return ans

ob = Solution()
nums = [2, 4, 5, 3, 1]
k = 2
print(ob.solve(nums, k))

입력

[2, 4, 5, 3, 1], 2

출력

6

복잡도 분석

이 알고리즘은 초기 합계 계산에 O(k), 이후 반복문에서 O(k)의 시간이 걸리므로 전체 시간 복잡도는 O(k)입니다. 모든 가능한 조합을 일일이 계산하는 브루트 포스 방식(O(2^k))보다 훨씬 효율적이며, 추가 공간 사용 없이 상수 공간 복잡도 O(1)로 문제를 해결할 수 있습니다.