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

파이썬으로 배열에서 요소를 삭제하며 얻을 수 있는 최대 포인트 구하기

N개의 요소로 이루어진 배열 A와 두 개의 정수 l, r이 주어져 있다고 가정해 보겠습니다(단, 1 ≤ ax ≤ 10^5, 1 ≤ l ≤ r ≤ N). 배열에서 임의의 요소 ax를 선택해 제거하면, 동시에 ax+1, ax+2 … ax+r에 해당하는 모든 요소와 ax−1, ax−2 … ax−l에 해당하는 모든 요소도 함께 배열에서 사라집니다. 이 작업을 수행하면 ax만큼의 포인트를 얻게 되며, 우리의 목표는 배열의 모든 요소를 제거한 후 획득한 총 포인트를 최대화하는 것입니다.

예를 들어 입력이 A = [2,4,3,10,5], l = 1, r = 2라고 하면, 출력은 18이 됩니다.

알고리즘 접근 방법

이 문제는 동적 프로그래밍(DP)을 활용해 효율적으로 해결할 수 있습니다. 단계별 절차는 다음과 같습니다.

  • n := 배열의 크기를 저장합니다.
  • max_val := 배열 내 최댓값을 구합니다.
  • count_list := 크기가 (max_val + 1)인 배열을 만들어 0으로 초기화한 뒤, 각 값의 등장 횟수를 기록합니다.
  • res := 크기가 (max_val + 1)인 DP 결과 배열을 만들어 0으로 초기화하고, res[0] = 0으로 설정합니다.
  • left := min(left, right)로 두 범위 값 중 작은 값을 사용합니다.
  • num을 1부터 max_val까지 반복하면서 다음을 수행합니다.
    • k := max(num − left − 1, 0)
    • res[num] := max(res[num − 1], num × count_list[num] + res[k])
  • 마지막으로 res[max_val]을 반환합니다.

핵심 아이디어는 각 숫자 num에 대해 '선택하지 않는 경우'와 '선택하는 경우'를 비교하는 것입니다. num을 선택하면 해당 값의 모든 등장 횟수만큼 포인트를 얻고, num±범위 안의 값들은 함께 제거되므로, num − left − 1까지의 최적 결과(res[k])만 이어서 활용할 수 있습니다. 이렇게 하면 불필요한 탐색 없이 O(max_val) 시간 복잡도로 최적해를 구할 수 있습니다.

예제 코드

아래의 파이썬 구현을 통해 더 자세히 이해해 보겠습니다.

def get_max_cost(array, left, right) :
    n = len(array)
    max_val = 0
    for i in range(n) :
        max_val = max(max_val, array[i])
    count_list = [0] * (max_val + 1)
    for i in range(n) :
        count_list[array[i]] += 1
    res = [0] * (max_val + 1)
    res[0] = 0
    left = min(left, right)
    for num in range(1, max_val + 1) :
        k = max(num - left - 1, 0)
        res[num] = max(res[num - 1], num * count_list[num] + res[k])
    return res[max_val]
array = [2,4,3,10,5]
left = 1
right = 2
print(get_max_cost(array, left, right))

입력

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

출력

18