백준 21909 Divisible by 3(Python)

서론

오랜만에 백준입니다. 수학을 많이 쓰려면 충분히 많이 쓸 수 있고 결국은 DP를 잘 활용해야 풀리는 문제입니다.

 

한 번 맛 봅시다.

 

본론

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

21909번: Divisible by 3

Divisible by 3

 

문제 설명

어떤 수열이 주어졌을 때 Weight를 정의합니다. Weight는 여기서 sum of pairwise product로 정의됩니다. 즉 모든 2개를 뽑는 조합의 경우의 값들을 모두 더했다는 겁니다.

 

그렇다면 하나의 수열이 주어졌을 때, 그 수열의 substring 중 weight가 3으로 나누어지는 건 몇 개나 될까? 가 질문입니다.

 

아이디어 1번

그런데 Pairwise product라고 하니 바로 누적합이 떠오릅니다. 각각 곱하는건 오래걸리지만 모두 더해놓고 한 번에 곱해버리면 쉽기 때문입니다. 이 때 자기 자신은 제외되어야 하니 제곱합의 누적합도 계산해주면 Pairwise product를 잘 구할 수 있겠습니다.

 

ex)

$$(a_1+a_2)^2 - ({a_1}^2+{a_2}^2)$$

 

이 방법을 사용하면 어떤 substring이 주어졌을 때 O(1)에 Weight를 구할 수 있습니다. 이 때 DP를 잘 써주면 답을 구할 수 있습니다.(있다는 사실만 알고 까먹음) 수열에 모두 mod 3을 씌워주고 3X3X(len(string))의 DP를 잘 채워주면 됩니다.

 

아이디어 2번

아이디어 1번에서 저기까지 생각하고 털렸기 때문에 2번을 떠올릴 수 밖에 없었습니다.

2번의 main 아이디어는 DP를 잘 정의하면 굳이 누적합을 쓸 필요도 없지 않을까? 였습니다.

 

예를 들어 string에서 i(i!=1)번째 element를 보고 있다고 해봅시다. 1부터 i-1번쨰 string에 i를 추가하면 weight가 어떻게 계산되어야 할까요?

만약 1부터 i-1번쨰 string까지의 weight가 주어졌다면, 새로운 weight는 (누적합 1~i-1) * i번째 element를 기존의 weight에 더하는 것과 같습니다. 새로 어떤 항들이 추가되어야할지 고민해보시면 답나옵니다.

 

그렇다면 DP로 쉽게 풀 수 있겠네요.

DP를 다음과 같이 정의합시다.

$$DP_{ijk}$$ = (i번째 string에 대해 weight의 3으로 나눈 나머지가 j이고 누적합을 3으로 나눈 나머지가 k인 substring의 개수

위를 쉽게 풀어쓰면 3x3행렬을 우리는 가질 건데

1 0 1

0 3 1

0 0 2

과 같이 있다면 1행 1열이 3이므로 weight의 나머지가 1이고 누적합의 나머지가 1인 개수가 3개 있다는 뜻입니다.

 

만약 그 다음(i+1)번째 숫자가 2라면 저 3을 보고 $$DP_{(i+1)00}$$에 값을 채워볼 수 있습니다. 왜냐면 1+2*1을 계산해주면 이는 새로운 weight가 될테고 (weight+new_val*prefixsum) 1에 2를 더해주고 3으로 나눈 나머지인 0이 새로운 누적합이 될 것이기 때문입니다.

 

위와 같은 과정으로 DP table을 하나씩 채워나가면 됩니다.

 

코드(python)

 

import sys

input = sys.stdin.readline
N = int(input())
arr = list(map(int,input().split()))
arr = list(map(lambda x : x%3, arr))

dp = [[[0,0,0],[0,0,0],[0,0,0]] for _ in range(N)]
dp[0][0][arr[0]] += 1

#i번째 원소를 무조건 포함하는 개수, 이 때 나머지가 0,1,2, 누적합이 0,1,2
ans = 1
for i in range(1,N):
    val = arr[i]
    dp[i][0][val] += 1
    for j in range(3):
        for k in range(3):
            dp[i][(j+k*val)%3][(k+val)%3] += dp[i-1][j][k]

    ans += sum(dp[i][0])
    
print(ans)

결론

아이디어 1번으로 풀고 싶었는데 아쉽게 실패했고 아이디어 2번과 아마 결이 같았을 것으로 예상되네요~

간만에 재밌는 문제였습니다. 누적합+DP