[Python] 28303번 자석

https://www.acmicpc.net/problem/28303

28303번: 자석

---

 

24/02/24

 

 

누적 합 응용으로, 아이디어를 떠올리니 쉽게 해결할 수 있었던 문제였다.

---

 

문제 접근 방식:

 

 

문제를 잘 읽고 분석을 해보았다.

 

배열을 $A$라고 하고, 배열의 $i$번째 원소를 $a_i$라고 한다면, 문제에서 요구하는 것은 $a_i - a_j - |i-j|K$의 값을 최대화시키는 것이다.

 

$N$의 최댓값이 $500,000$이기 때문에, $\mathcal{O}(N^2)$의 시간 복잡도로 완전 탐색을 진행한다면 당연히 시간초과가 난다.

 

때문에 $\mathcal{O}(N)$의 시간 복잡도로 구현할 필요가 있다.

 

처음에는 누적 합을 이리저리 사용하면 되지 않을까 생각했지만, 구체적인 방법은 떠오르지 않아 잠시 미뤄두었다.

 

또 다른 방법으로는 DP가 생각났다.

 

$DP[i]$를 $i$번째 항을 볼 때 문제에서 요구하는 최댓값이라고 설정했다.

 

하지만 이래도 뚜렷한 점화식이 보이지 않아 포기하게 되었다.

 

그러던 중 텔레스코핑의 아이디어가 생각이 났다.

 

즉, $a_i - a_j$는 $a_i - a_{i-1} + a_{i-1} - a_{i-2} + \cdots + a_{j+1} - a_j$와 같은 형태로 변형될 수 있음을 깨달았다.

 

문제에서 요구하는 에너지의 변화량의 형태는 $a_i - a_j - |i-j|K$이므로, $a_i - a_{i-1}$들을 저장하는 것이 아니라, $a_i - a_{i-1} - K$들을 저장하고, 그들의 구간 합을 구한다면 문제에서 요구하는 에너지의 변화량을 쉽게 구할 수 있음을 알았다.

 

구간 합을 빠르게 계산하는 방법은 누적 합이다.

 

누적 합 배열의 두 원소의 차를 통해 구간 합을 빠르게 계산할 수 있다. 편의 상 이 누적 합 배열을 $B$라고 하자. 

 

즉, 누적 합으로 원본 배열을 잘 변형시켜서, $a_i - a_j - |i-j|K$의 문제를 $b_i - b_j$의 최댓값을 찾는 문제로 변형시킬 수 있었다.

 

이는 $\mathcal{O}(N)$으로 해결할 수 있는데, 배열 $B$의 원소를 순서대로 보다가, 현재 배열의 원소 $b_i$와 이전까지 봤던 원소들의 최솟값 $m$을 빼준 값을 구했다.

 

그 값이 지금까지 봤던 $b_i - b_j$의 최댓값보다 크다면 그 최댓값을 갱신시켜 주었다.

 

이후, $m$보다 $b_i$의 값이 더 작다면 $m$값을 갱신해 주었다.

 

N극을 왼쪽에 두냐 오른쪽에 두냐에 따라서도 값이 다르게 나올 수 있으므로, 이러한 작업을 두 번 거쳐서 둘 중에 더 큰 값을 출력하도록 해주었다.

---

아래는 내가 위의 접근 방식과 같이 작성한 파이썬 코드이다. 더 보기를 누르면 확인할 수 있다.

# 28303번 자석
# 누적 합, DP
'''
접근 방법:
에너지 변화량
ai - aj - abs(i-j)*K
누적합으로 잘 변형
ai - aj (i > j)의 최댓값을 찾는 문제로 변형
DP[i] = i번째 항을 볼 때 ai - aj의 최댓값
'''
import sys
input = sys.stdin.readline

N, K = map(int, input().rstrip().split())
A = list(map(int, input().rstrip().split()))
diff1, diff2 = [], []
for i in range(N-1):
    diff1.append(A[i+1] - A[i] - K)
    diff2.append(A[i] - A[i+1] - K)
diff_accum_sum1, diff_accum_sum2 = [0], [0]
for i in range(N-1):
    diff_accum_sum1.append(diff1[i] + diff_accum_sum1[-1])
    diff_accum_sum2.append(diff2[i] + diff_accum_sum2[-1])
DP1, DP2 = [diff_accum_sum1[1]], [diff_accum_sum2[1]]
min1, min2 = min(0, diff_accum_sum1[1]), min(0, diff_accum_sum2[1])
for i in range(2, N):
    DP1.append(max(DP1[-1], diff_accum_sum1[i] - min1))
    DP2.append(max(DP2[-1], diff_accum_sum2[i] - min2))
    if diff_accum_sum1[i] < min1:
        min1 = diff_accum_sum1[i]
    if diff_accum_sum2[i] < min2:
        min2 = diff_accum_sum2[i]
print(max(DP1[-1], DP2[-1]))