세그먼트 트리 문제 풀이 모음 220924
백만년만에 문제 풀이 정리를 한다.
물론 시간 날 때마다 문제는 꾸준히 풀어오고 있다.
그동안 정리 안 한 것은 브론즈 문제들을 굳이 풀이를 정리하기에는 불필요하다고 느껴서였고, 가끔 약간의 메모가 필요한 문제들이 있으면 정리해보려고 한다.
백준 BOJ 8120번 : Coding of Permutations
문제 난이도 : Gold I
알고리즘 분류 : 세그먼트 트리
길이 N인 수열로 a_i = i번째 위치의 왼쪽에 있는 더 큰 수의 개수가 주어질 때, N에 대한 순열을 찾는 문제이다.
예를 들어 (0, 0, 1, 0, 2, 0, 4)가 주어진다면 (1, 5, 2, 6, 4, 7, 3)이 답인데 왜냐하면 1의 왼쪽에 더 큰 수가 0개, 5의 왼쪽에 더 큰 수가 0개, 2의 왼쪽에 더 큰 수가 1개, 6의 왼쪽에 더 큰 수가 0개, ...로 규칙을 만족하기 때문이다.
이 문제의 경우에는 역순으로 체크해보면서 답을 구하면 된다.
예를 들어 4가 나왔다면, 1 ~ 7 중에서 4개의 더 큰 수를 갖는 3이 맨 오른쪽에 위치함을 알 수 있다.
이제 1 ~ 7의 목록에서 3을 지운다. (즉, {1, 2, 4, 5, 6, 7})
이제 오른쪽에서 두 번째에 있는 수가 0이므로, 남은 수 중에서 0개의 더 큰 수를 갖는 7이 오른쪽에서 두 번째에 옴을 알 수 있다.
이런 식으로 풀어주면 되는데, 목록에서 수를 제거해야 하기 때문에 세그먼트 트리를 사용한다.
먼저 트리에 1 ~ N의 위치에 수를 1씩 넣어주고, 순열 (0, 0, 1, 0, 2, 0, 4)를 벡터 u에 저장했다고 가정하자.
세그먼트 트리는 벡터 v에 구현되었다고 가정하자.
그러면 v[1](= 루트 노드, 즉 트리의 현재 사이즈) - u[i]번째로 작은 원소를 출력하고 그것을 지워주면 된다.
만약 u[i] ≥ v[1]인 순간이 한 번이라도 있으면 그러한 순열은 존재하지 않는다.
#include <bits/stdc++.h>
#define int long long
using namespace std;
int N, Max = 3e4 + 1;
vector<int> v(Max * 4);
void f(int n, int b, int e, int idx) {
if(idx < b || e < idx) return;
v[n]++;
if(b == e) return;
f(n*2, b, (b+e)/2, idx);
f(n*2 + 1, (b+e)/2 + 1, e, idx);
}
int g(int n, int b, int e, int cnt) {
v[n]--;
if(b == e) return b;
if(cnt <= v[n*2]) return g(n*2, b, (b+e)/2, cnt);
else return g(n*2 + 1, (b+e)/2 + 1, e, cnt - v[n*2]);
}
main() {
ios_base::sync_with_stdio(false);
cin.tie(NULL), cout.tie(NULL);
cin >> N;
for(int i=1; i<=N; i++) f(1, 1, Max, i);
vector<int> u(N+1);
for(int i=1; i<=N; i++) cin >> u[i];
vector<int> ans(N+1);
for(int i=N; i>=1; i--) {
if(u[i] >= v[1]) {
cout << "NIE\n";
return 0;
}
ans[i] = g(1, 1, Max, v[1] - u[i]);
}
for(int i=1; i<=N; i++) cout << ans[i] << "\n";
}
백준 BOJ 15459번 : Haybale Feast
문제 난이도 : Gold I
알고리즘 분류 : 세그먼트 트리, 두 포인터
길이 N인 수열 A, B에 대해 연속한 부분 수열을 골라 a_l ~ a_r 값이 M 이상인 것들 중 b_l ~ b_r 중 최댓값이 최소가 되도록하는 값을 구하는 문제이다.
수열 B에 대해 구간의 최솟값을 다루는 세그먼트 트리를 만들고, 수열 A에 대해 투 포인터 알고리즘을 써주면 된다.
항상 느끼지만 USACO에 정석적이고 좋은 문제가 많은 것 같다.
#include <bits/stdc++.h>
#define int long long
using namespace std;
vector<int> v, w;
int init(int n, int b, int e) {
if(b == e) return v[n] = w[b];
int lv = init(n*2, b, (b+e)/2);
int rv = init(n*2 + 1, (b+e)/2 + 1, e);
return v[n] = max(lv, rv);
}
int f(int n, int b, int e, int l, int r) {
if(r < b || e < l) return INT_MIN;
if(l <= b && e <= r) return v[n];
int lv = f(n*2, b, (b+e)/2, l, r);
int rv = f(n*2 + 1, (b+e)/2 + 1, e, l, r);
return max(lv, rv);
}
main() {
ios_base::sync_with_stdio(false);
cin.tie(NULL), cout.tie(NULL);
int N, M; cin >> N >> M;
vector<int> u(N+1);
w.resize(N+1);
for(int i=1; i<=N; i++) {
int x; cin >> x >> w[i];
u[i] = u[i-1] + x;
}
v.resize(N*4);
init(1, 1, N);
int i = 1, j = 1, ans = INT_MAX;
while(j <= N) {
int sum = u[j] - u[i-1];
if(sum < M) j++;
else {
ans = min(ans, f(1, 1, N, i, j));
i++;
if(i > j) j = i;
}
}
cout << ans << "\n";
}
백준 BOJ 12803번 : Peter and the Textbook
문제 난이도 : Silver II
알고리즘 분류 : 세그먼트 트리
1 ~ N 페이지로 구성된 책이 있고 x번째 페이지를 제거한다면 같이 제본된 N-x+1번째 페이지도 같이 제거된다고 할 때, 두 가지의 쿼리를 처리하는 문제이다.
1. - x : x번째 페이지와 N-x+1번째 페이지를 제거한다.
2. ? x : 현재 남은 페이지 중에 앞에서부터 x번째 페이지의 번호를 출력한다.
N이 매우 작으므로 세그먼트 트리를 쓸 필요는 없지만, 세그먼트 트리의 연습을 위해 세그트리로 문제를 풀어보자.
먼저 1 ~ N번째 리프 노드에 대해 모두 1로 채워주고, 그 다음 - 쿼리에 대해서는 루트 노드에서부터 1씩 감소시켜가면서 리프 노드까지 값을 처리해주면 된다.
그리고 ? 쿼리에 대해서는 왼쪽 자식의 값이 cnt 값보다 큰지 작은지를 판단하여 내려가는 방향을 달리해주어 탐색해주면 된다.
#include <bits/stdc++.h>
#define int long long
using namespace std;
int Max = 1e2 + 1;
vector<int> v;
void f(int n, int b, int e, int idx) {
if(idx < b || e < idx) return;
v[n]++;
if(b == e) return;
f(n*2, b, (b+e)/2, idx);
f(n*2 + 1, (b+e)/2 + 1, e, idx);
}
void g(int n, int b, int e, int idx) {
if(idx < b || e < idx) return;
v[n]--;
if(b == e) return;
g(n*2, b, (b+e)/2, idx);
g(n*2 + 1, (b+e)/2 + 1, e, idx);
}
int h(int n, int b, int e, int cnt) {
if(b == e) return b;
if(cnt <= v[n*2]) return h(n*2, b, (b+e)/2, cnt);
else return h(n*2 + 1, (b+e)/2 + 1, e, cnt - v[n*2]);
}
main() {
ios_base::sync_with_stdio(false);
cin.tie(NULL), cout.tie(NULL);
int T; cin >> T;
while(T--) {
v.clear();
v.resize(Max*4);
int N, M; cin >> N >> M;
for(int i=1; i<=N; i++) f(1, 1, Max, i);
while(M--) {
char Q; cin >> Q;
int x; cin >> x;
if(Q == '-') {
g(1, 1, Max, x);
g(1, 1, Max, N-x+1);
}
else if(Q == '?') cout << h(1, 1, Max, x) << "\n";
}
}
}
백준 BOJ 2104번 : 부분배열 고르기
문제 난이도 : Platinum V
알고리즘 분류 : 세그먼트 트리
길이 N인 수열이 주어질 때 연속하는 부분 수열을 골라 그 점수가 (v[l] + v[l+1] + ... + v[r]) * min(v[l] ~ v[r])이라고 할 때 점수의 최댓값을 구하는 문제이다.
백준 BOJ 6549번 : 히스토그램에서 가장 큰 직사각형 문제를 떠올려보면 같은 방법으로 풀이가 가능하다는 것을 알 수 있다.
다만 여기서는 직사각형의 넓이가 아닌 다른 점수를 사용하므로, v[l] + ... + v[r] 값을 누적 합 배열을 사용하여 w[r] - w[l-1]로 O(1)에 구할 수 있고, 이를 활용하여 점수 값을 비교하여 최댓값을 구해주면 된다.
#include <bits/stdc++.h>
#define int long long
using namespace std;
int N;
vector<int> v, u, w;
int init(int n, int b, int e) {
if(b == e) return u[n] = b;
int lv = init(n*2, b, (b+e)/2);
int rv = init(n*2 + 1, (b+e)/2 + 1, e);
if(v[lv] < v[rv]) return u[n] = lv;
else return u[n] = rv;
}
int g(int n, int b, int e, int l, int r) {
if(r < b || e < l) return -1;
if(l <= b && e <= r) return u[n];
int lv = g(n*2, b, (b+e)/2, l, r);
int rv = g(n*2 + 1, (b+e)/2 + 1, e, l, r);
if(lv < 0) return rv;
if(rv < 0) return lv;
if(v[lv] < v[rv]) return lv;
else return rv;
}
int f(int l, int r) {
int idx = g(1, 1, N, l, r);
int ret = (w[r] - w[l-1]) * v[idx];
if(idx+1 <= r) ret = max(ret, f(idx+1, r));
if(idx-1 >= l) ret = max(ret, f(l, idx-1));
return ret;
}
main() {
ios_base::sync_with_stdio(false);
cin.tie(NULL), cout.tie(NULL);
cin >> N;
v.resize(N+1);
w.resize(N+1);
for(int i=1; i<=N; i++) {
cin >> v[i];
w[i] = w[i-1] + v[i];
}
u.resize(N*4);
init(1, 1, N);
cout << f(1, N) << "\n";
}
백준 BOJ 1989번 : 부분배열 고르기 2
문제 난이도 : Platinum V
알고리즘 분류 : 세그먼트 트리
위의 문제와 동일하나 이번에는 구간의 주소까지 구해주어야 하는 문제이다.
시간 복잡도를 줄여야 하거나 다른 전략을 사용해야 하는 것은 아니고, left 값과 right 값을 인자로 같이 전달해주기만 하면 된다.
구조체를 사용하여 3개의 값을 한꺼번에 전달할 수 있도록 구현하였다.
#include <bits/stdc++.h>
#define int long long
using namespace std;
int N;
vector<int> v, u, w;
int init(int n, int b, int e) {
if(b == e) return u[n] = b;
int lv = init(n*2, b, (b+e)/2);
int rv = init(n*2 + 1, (b+e)/2 + 1, e);
if(v[lv] < v[rv]) return u[n] = lv;
else return u[n] = rv;
}
int g(int n, int b, int e, int l, int r) {
if(r < b || e < l) return -1;
if(l <= b && e <= r) return u[n];
int lv = g(n*2, b, (b+e)/2, l, r);
int rv = g(n*2 + 1, (b+e)/2 + 1, e, l, r);
if(lv < 0) return rv;
if(rv < 0) return lv;
if(v[lv] < v[rv]) return lv;
else return rv;
}
struct s { int l, r, a; };
s f(int l, int r) {
int idx = g(1, 1, N, l, r);
int area = (w[r] - w[l-1]) * v[idx];
s ret = {l, r, area};
if(idx+1 <= r) {
s tmp = f(idx+1, r);
if(tmp.a > ret.a) ret = tmp;
}
if(idx-1 >= l) {
s tmp = f(l, idx-1);
if(tmp.a > ret.a) ret = tmp;
}
return ret;
}
main() {
ios_base::sync_with_stdio(false);
cin.tie(NULL), cout.tie(NULL);
cin >> N;
v.resize(N+1);
w.resize(N+1);
for(int i=1; i<=N; i++) {
cin >> v[i];
w[i] = w[i-1] + v[i];
}
u.resize(N*4);
init(1, 1, N);
s ans = f(1, N);
cout << ans.a << "\n";
cout << ans.l << " " << ans.r << "\n";
}
백준 BOJ 11143번 : Beads
문제 난이도 : Gold I
알고리즘 분류 : 세그먼트 트리
N개의 상자가 있고, 쿼리 P에 대해서는 a번 상자에 구슬 b개를 추가하고, 쿼리 Q에 대해서는 a번 상자부터 b번 상자까지 모든 구슬의 개수의 합을 구하는 처리를 하는 문제이다.
구간 합 구하기 문제를 활용하여 거의 비슷하게 풀이할 수 있다.
일단 초기 상자들에는 구슬이 모두 0개 들어있으므로 init 함수와 같은 것을 구현할 필요가 없다.
값을 갱신할 때는 특정 인덱스의 값을 바꾸는 것이 아닌 더하는 것이므로, diff = b로 설정해주고 v[a] += b와 같이 해주면 갱신이 가능하다.
#include <bits/stdc++.h>
#define int long long
using namespace std;
vector<int> v, u;
void upd(int n, int b, int e, int idx, int diff) {
if(idx < b || e < idx) return;
u[n] += diff;
if(b < e) {
upd(n*2, b, (b+e)/2, idx, diff);
upd(n*2 + 1, (b+e)/2 + 1, e, idx, diff);
}
}
int f(int n, int b, int e, int l, int r) {
if(r < b || e < l) return 0;
if(l <= b && e <= r) return u[n];
return f(n*2, b, (b+e)/2, l, r) + f(n*2 + 1, (b+e)/2 + 1, e, l, r);
}
main() {
ios_base::sync_with_stdio(false);
cin.tie(NULL), cout.tie(NULL);
int T; cin >> T;
while(T--) {
int N, M, K; cin >> N >> M >> K;
v.clear();
v.resize(N+1);
u.clear();
u.resize(N*4);
M += K;
while(M--) {
char Q; int a, b; cin >> Q >> a >> b;
if(Q == 'P') {
int diff = b;
v[a] += b;
upd(1, 1, N, a, diff);
}
else if(Q == 'Q') cout << f(1, 1, N, a, b) << "\n";
}
}
}