Prefix Sum
Prefix sum이란 해당 인덱스까지의 value 들의 누적 합을 말한다.
| Index | 0 | 1 | 2 | 3 | 4 | 5 | 6 |
|---|---|---|---|---|---|---|---|
| Value | 3 | 5 | -2 | -1 | 4 | -10 | |
| Prefix Sum | 0 | 3 | 8 | 6 | 5 | 9 | -1 |
위에서 예를 들면 1번 인덱스까지의 누적합은 3이니까 prefix sum이 3이고, 5번 인덱스까지의 누적합은 9이므로 prefix sum은 9이다. 이걸 이용하면 index 1, 5의 prefix sum을 알 때 index 2~4의 value의 합을 구할 수 있다.
Prefix sum을 활용하는 문제를 풀어보자.
Codeforces 466C: Number of Ways
배열의 모든 element를 세 개의 연속된 부분으로 분할하고 각 element의 합이 같도록 하는 방법의 수를 구하는 문제. 즉 sigma(1, i-1, ak) = sigma(i, j, ak) = sigma(j+1, n)을 만족시키는 (i, j)의 순서쌍의 개수를 구해야 한다. 이때 2 ≤ i ≤ j ≤ n - 1이므로 각
처음에 내가 짠 코드. N의 최댓값이 500만이기 때문에 딱 봐도 시간초과가 날 것 같지만 그래도 재귀로 짜봤다. prefix sum을 활용할 생각은 하지도 않고 신나서 풀음..
#include <iostream>
#include <vector>
using namespace std;
long long N, total, cnt; // N의 범위가 10^5이므로 전체 방법의 수가 int의 범위를 넘어갈 수 있음
vector<int> v;
void partition(int start, int sum, int parts) {
int tmp, i;
if (parts == 1) { // 세 덩어리로 partitioning이 완료되었으므로 cnt 증가시킴
cnt++;
return;
}
tmp = 0;
for (i = start; i < N - (parts - 1); i++) {
tmp += v[i];
if (tmp == sum) partition(i + 1, sum, parts-1);
}
}
int main() {
ios_base::sync_with_stdio(false);
cin.tie(0);
int i, tmp;
cin >> N;
for (i = 0; i < N; i++) {
cin >> tmp;
v.push_back(tmp);
total += tmp;
}
if (total % 3 != 0) cout << 0 << endl; // 전체 합이 3의 배수가 아니면 세 개의 subarray로 partitioning 할 수 없다.
else {
partition(0, total / 3, 3);
cout << cnt << endl;
}
return 0;
}
친절한 코드포스는 백준과는 달리 테케를 모두 공개하므로 N이 백만일 때 시간초과가 발생한다는 것을 알게 되었다.
교수님의 코드를 참고해 prefix sum을 이용하여 해당 인덱스까지의 누적합을 prefix[]에 저장하는 방식으로 작성한 코드. 각 element의 범위가 백만 이하이므로 p를 int가 아닌 long long 타입의 벡터로 선언해야 함.
#include <iostream>
#include <vector>
#define MAX 500001
using namespace std;
int main() {
ios_base::sync_with_stdio(false);
cin.tie(0);
long long N, total, cnt, cnt3, i;
vector<long long> v(MAX); // value
vector<long long> p(MAX); // prefix sum
p[0] = 0;
v[0] = 0;
cin >> N;
for (i = 1; i <= N; i++) {
cin >> v[i];
p[i] = p[i - 1] + v[i];
}
total = p[N];
if (total % 3 != 0) cout << 0 << endl;
else {
cnt = 0;
cnt3 = 0;
if (p[1] == total / 3) cnt3++;
for (i = 2; i <= N - 1; i++) { // 마지막 partition에도 최소 1개의 element가 있어야 하기 때문에
if (p[i] == (total / 3) * 2) cnt += cnt3; // 현재 보고있는게 전체의 2/3이면 지금껏 본 1/3의 개수를 더한다.
if (p[i] == total / 3) cnt3++;
}
cout << cnt << endl;
}
return 0;
}
Codeforces 1003C: Intense Heat
prefix sum을 안 썼으면 매번 sum을 구하느라 O(N)이 더 곱해졌을거다. 이래서 많이 아는게 중요함.
오랜만에 소숫점 계산했다..
#include <iostream>
#include <vector>
#include <algorithm> // max
#define MAX 5001
using namespace std;
int main() {
ios_base::sync_with_stdio(false);
cin.tie(0);
int N, K, i, j;
double answer, tmp;
vector<int> v(MAX);
vector<int> p(MAX);
cin >> N >> K;
p[0] = v[0] = 0;
for (i = 1; i <= N; i++) {
cin >> v[i];
p[i] = p[i - 1] + v[i];
}
answer = 0.0;
for (i = K; i <= N; i++) { // not less than K consecutive days
for (j = 1; j <= N - (i - 1); j++) {
tmp = ((double)p[j + i - 1] - p[j - 1]) / (double)i;
answer = max(tmp, answer);
}
}
printf("%.15f\n", answer);
return 0;
}
Codeforces 1285B: Just Eat It!
쉽네~하며 풀었다가 시간초과난 문제. O(N^2)에 풀었는데 틀린 테케를 보니 O(N)에 풀어야 하는 것 같다.
이렇게 풀면 시간초과 뜸.
#include <iostream>
#include <vector>
#include <algorithm> // max, min
#define MAX 100001
using namespace std;
int main() {
ios_base::sync_with_stdio(false);
cin.tie(0);
int T, N, i, j;
long long total, partial, tmp;
cin >> T;
while (T--) {
cin >> N;
vector<int> v(MAX);
vector<long long> p(MAX);
partial = 0;
p[0] = 0;
for (i = 1; i <= N; i++) {
cin >> v[i];
p[i] = p[i - 1] + v[i];
}
total = p[N];
for (i = 1; i <= N - 1; i++) { // size of the segment
for (j = 1; j <= N - (i - 1); j++) {
tmp = p[j + i - 1] - p[j - 1];
partial = max(partial, tmp);
}
}
if (total > partial) cout << "YES" << endl;
else cout << "NO" << endl;
}
return 0;
}
v의 앞뒤 일부를 지워 total 이상으로 만드는 방법이 존재한다면 NO를 출력하면 된다.
교수님의 코드는 아직도 이해하지 못했지만 나름 풀어봤다.
예를 들어 v[] = {1, -5, 4, -8, 7}의 배열이 있다고 하자.
prefix sum과 suffix sum(뒤에서부터 거꾸로인 prefix sum이라 생각하면 됨)을 구해보면 다음과 같다.
| Index | 0 | 1 | 2 | 3 | 4 | 5 |
|---|---|---|---|---|---|---|
| Value | 1 | -5 | 4 | -8 | 7 | |
| Prefix Sum | 1 | -4 | 0 | -8 | -1 | |
| Suffix Sum | -1 | -2 | 3 | -1 | 7 |
이제 마치 퀵소트를 할 때 양 끝에서부터 움직이는 것처럼 prefix sum과 suffix sum 배열의 값을 양 끝에서부터 관찰하자.
i가 2일 때
i=2일 때 prefix sum값이 -4로 0보다 작다. prefix sum[2]은 v[1]과 v[2]의 합이고, 이는 v[1] + v[2]를 더한 값은 전체 배열의 합을 오히려 작게 만든다는 것을 의미한다.
실제로 v에서 1~5번 인덱스를 다 더한 값은 1+(-5)+4+(-8)+7 = -1이지만, 1번과 2번 인덱스를 제외하고 더한 값은 4+(-8)+7 = 3이다. 그렇기 때문에 우리는 prefix sum을 사용해 차라리 없는게 이득인 부분, 즉 prefix sum이 양수가 아닐 때가 존재하는지를 찾는다.
만약 prefix sum이 0보다 같거나 작은 지점이 하나라도 있다고 하자. 그러면 Adel은 이 부분을 버릴 것이고 그러면 Yasser의 tastiness(total값)보다 크거나 같은 tastiness 값을 가지게 되므로 프로그램은 NO를 출력할 것이다.
마찬가지로 suffix sum을 활용해 동일한 루틴을 v의 맨 끝에서부터 수행한다.
이렇게 해도 시간초과가 뜨는데, prefix sum나 suffix sum에서 0 이하의 값을 발견하자마자 NO를 출력하도록 코드를 작성하자. flag 값이 한 번 false가 되었으면 나머지는 더 이상 비교할 필요가 없다.
#include <iostream>
#include <vector>
#include <algorithm>
#define MAX 100002
using namespace std;
int N;
vector<int> v(MAX);
vector<long long> p(MAX); // prefix sum
vector<long long> s(MAX); // suffix sum
// key: 앞에서부터랑 뒤에서부터 각각 더해봄. 그 값이 0보다 작으면 NO임! 왜냐하면 그 값이 없었더라면 total보다 큰 값이 나왔을 테니깐~
bool check() {
int i;
if (v[1] <= 0 || v[N] <= 0) return false;
for (i = 1; i <= N; i++) {
if (p[i] <= 0 || s[N - i + 1] <= 0) return false;
}
return true;
}
int main() {
ios_base::sync_with_stdio(false);
cin.tie(0);
int T, i;
cin >> T;
while (T--) {
cin >> N;
p[0] = v[0] = 0;
for (i = 1; i <= N; i++) {
cin >> v[i];
p[i] = p[i - 1] + v[i];
}
s[N + 1] = 0;
for (i = N; i >= 1; i--) s[i] = s[i + 1] + v[i];
if (check()) printf("YES\n");
else printf("NO\n");
}
return 0;
}