배열에서 일부를 골라 공차가 매우 큰 등차수열이 되도록 만드는 방법의 수
E - Count Arithmetic Subsequences (atcoder.jp)
배열 A에서 일부를 골라 길이가 i = 1,2,3,..,n인 등차수열을 만드는 방법의 수를 각각 구하는 문제
dp[i][j][d]을 길이가 j이고 공차가 d이며 등차수열의 끝 항이 A[i]인 등차수열을 만드는 방법의 수라고 해보자.
길이가 j-1이고 끝항이 A[k], k = 0,1,2,..,i-1인 모든 등차수열을 조사하여 A[i] - A[k] = d인 등차수열을 발견한다면,
그러한 등차수열에 A[i]를 추가하기만 하면 되므로 dp[i][j][d] += dp[k][j-1][d]가 된다
for j in range(1,n+1): #길이
for i in range(n): #끝 항
for k in range(i): #0~i-1항
d = A[i] - A[k]
dp[i][j][d] += dp[k][j-1][d]
이러면 $O(N^{3})$에 문제를 해결할 수 있다
여기서 문제는 배열의 원소 A가 최대 $10^{9}$이므로, 공차도 최악의 경우 $10^{9} - 1$까지 가능하다
dp[i][j][d]에서 3번째가 공차인데 $10^{9}$가 들어가면 메모리 초과인데
{} dictionary로 메모리를 절약할 수 있다
이러면 공차는 아무리 많아야 $O(N)$개이므로 dp가 차지하는 공간은 $O(N^{3})$
다음으로 문제는 dp배열 초기화이다.
먼저 길이가 1인 경우에는? 공차가 없는데 어떡하지?
사실 길이가 1인 등차수열은 A[0], A[1], A[2],...,A[n-1] 각각으로 총 n개인것은 명확하다.
얘들은 공차가 없으니까 신경쓰지말고 길이가 2부터 시작하면 충분하다.
길이가 2인 경우를 $O(N^{2})$으로 초기화시킨다
dp = [[{} for _ in range(n+1)] for _ in range(n+1)]
for i in range(n): #끝 항
for j in range(i): #0~i-1항
d = A[i] - A[j] #공차
#끝항이 A[i], 길이가 2, 공차가 d
dp[i][2][d] = dp[i][2].get(d,0) + 1
그러면 길이가 3인 경우부터 dp배열을 채워나가면 된다
mod = 998244353
for i in range(3,n+1): #길이
for j in range(n): #끝항
for k in range(j): #0~j-1항
d = A[j] - A[k] #공차
#끝항이 A[j]이고 길이가 i이며 공차가 d인 등차수열은..
#끝항이 A[k]이고 길이가 i-1이며 공차가 d인 등차수열이 존재한다면..
dp[j][i][d] = dp[j][i].get(d,0) + dp[k][i-1].get(d,0)
dp[j][i][d] %= mod
이제 모든 dp배열을 찾았으므로, 길이가 i = 1,2,3,..,n에 대하여 등차수열을 만드는 모든 방법의 수를 구하면 된다
sum(dp[j][i])가 길이 i인 등차수열을 만드는 방법의 수일 것이다
result = [n] #길이가 1인 경우는 n개
for i in range(2,n+1): #길이가 2부터
v = 0
for j in range(n):
for d in dp[j][i]: #끝항이 A[j]이고 길이가 i인 dp[j][i]의 모든 공차 d에 대하여
v += dp[j][i][d]
v %= mod
result.append(v)
'알고리즘 > 다이나믹 프로그래밍' 카테고리의 다른 글
홀수 길이의 연속 부분 배열의 합 중 가장 큰 값을 구하는 방법 (0) | 2024.08.05 |
---|---|
index와 value를 바꾸는 다이나믹 프로그래밍 트릭 (0) | 2024.07.29 |
다이나믹 프로그래밍을 이용한 partition minimum sum of difference of two subset(합의 차이가 최소가 되는 분할) 문제 해법 (0) | 2024.07.08 |
가방이 여러개인 multiple 0-1 knapsack(배낭 문제) 해법 배우기 (0) | 2024.07.03 |
0-1 knapsack 배낭 문제 해법 반드시 알아야하는 디테일(물건을 중복해서 담는 경우) (0) | 2024.07.03 |