배열에서 일부를 골라 공차가 매우 큰 등차수열이 되도록 만드는 방법의 수

E - Count Arithmetic Subsequences (atcoder.jp)

E - Count Arithmetic Subsequences

 

 

배열 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)