취미로 코딩하기

※ 글쓴이는 취미로 코딩을 익혀보는 사람이라 정확하지 않은 내용을 담고 있을 수 있다 ※

이번에 볼 문제는 백준 3673번 문제인 나눌 수 있는 부분 수열이다.
문제는 아래 링크를 확인하자.

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

모든 부분합을 다 살피는 것은 O(min(dN, N^2))에 해낼 수 있는데, 이는 이 문제를 해결하기에는 부족한 시간복잡도이다.

연속하는 부분 수열의 합에 대한 것을 알고싶은 문제이므로, prefix sum을 구해 문제를 풀어보자.

i번째 항부터 j번째 항까지의 합이 d로 나누어떨어진다는 말은 초항부터 i-1번째 항까지의 합과 초항부터 j번째 항까지의 합의 차가 d의 배수라는 뜻이므로, 둘을 각각 d로 나눈 나머지는 같다는 것을 알 수 있다.

따라서, prefix sum의 d로 나눈 나머지를 구해나가면서, d로 나눈 나머지가 같은 숫자들의 쌍의 개수를 세는 것으로 문제를 해결할 수 있다.

단, 초항부터 i번항까지의 합이 d의 배수라면 그 구간합 자체도 답의 하나이므로 다른 경우와 다르게 처리해야한다는 점에 유의하자.

아래는 제출한 소스코드이다.

#include <iostream>
#include <cstring>
using namespace std;
typedef long long ll;

int cnt[1000000];

int main() {
	ios::sync_with_stdio(0);
	cin.tie(0);

	int T; cin >> T;
	while (T--) {
		memset(cnt, 0, sizeof(cnt));
		cnt[0]++;
		ll ans = 0;
		int mod, N; cin >> mod >> N;
		int total = 0;
		while (N--) {
			int x; cin >> x;
			total += x;
			total %= mod;
			ans += cnt[total]++;
		}
		cout << ans << '\n';
	}
}