백준 21909 / C++

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

21909번: Divisible by 3

 

[ 풀이 ] 

문제의 정의에 나온 weight는 식으로 정리하면,  1/2*{(a[l]+....+a[r])^2 - (a[l]^2 + .... +a[r]^2)}이 된다. 

psum[i]=a[1]+....+a[i]라고 하고, T[i]=a[1]^2+....+a[i]^2이라 하자. 

(psum[r]-psum[l-1])^2==T[r]-T[l-1] (mod3) 인 경우의 수를 구하면 된다. 

이걸 O(N)에 잘 처리해보자. 아이디어는 스위핑을 하면서 psum[i]과 T[i]의 나머지를 저장하는 것이다. 

어차피 mod3으로 보기 때문에 dp[x][y]:= psum[i]=x (mod3) 이고 T[i]=y (mod3) 인 것의 개수를 갱신하는 것이다. 

이제 dp[x][y]가 갱신되었다고 할 때, 그 다음 값을 보면서 (psum-x)^2==T-y (mod3)인 0<x,y<3를 찾는다.

만족한다면 dp[x][y]를 더해주면 된다. 이걸 1번 스위핑하면 답이 된다. 

어려운 문제였다. 

 

[ Code ] 

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
ll n, k, dp[3][3];

int main() {
	ios::sync_with_stdio(0); cin.tie(0);
	cin >> n;
	ll S = 0, T = 0, ans = 0;
	dp[0][0] = 1;
	for (int i = 1; i <= n; i++) {
		cin >> k; k %= 3;
		S += k; // S[i]
		T += k * k; // T[i]
		for (int x = 0; x < 3; x++) {
			for (int y = 0; y < 3; y++) {
				// (S-x)^2=T-y (mod3) 
				if (((S - x) * (S - x)) % 3 == (T - y) % 3) {
					ans += dp[x][y];
				}
			}
		}
		dp[S % 3][T % 3]++; // 갱신
	}
	cout << ans << '\n';
}