777E Hanoi Factory

문제

• 링크 - http://codeforces.com/contest/777/problem/E

문제 분류

• 난이도 - 중
• 태그 - dynamic programming, data structure, segment tree

문제 풀이

내부 반지름의 길이 \(a_i\), 외부 반지름의 길이 \(b_i\), 그리고 높이 \(h_i\)가 주어진 여러 개의 링들을 최대한 높게 쌓는 문제다. 일단 아래와 같이 2차원 평면으로 보면 (1) \(b_i\)가 증가하는 순으로, (2) \(b_i\)가 증가하는 순으로 정렬해야할 필요성을 느낄 수 있다.

완전 탐색을 할 경우 \(O(2^n)\)으로 문제를 해결할 수 있다. 너무 느리므로 동적계획법을 고민해보면 아래와 같은 점화식을 찾을 수 있다.

\(dp[i]\): \(i\)번 링을 맨 밑으로 하여 쌓을 수 있는 최대 높이 (단, 위에서 언급한대로 정렬되어 있어야 한다)

\(dp[i]=\underset {j < i \ and \ a_i < b_j}{max}(dp[j]) + h_i\)

위 점화식을 반복문을 이용하여 계산하면, \(O(n^2)\)에 문제를 해결할 수 있지만, 역시 정답을 내기에는 부족하다. 하지만 잘보면 위 식에서 \(dp[i]\)를 계산할 때, 이미 계산된 값 \(dp[j] \ (j < i)\) 중 \(a_i < b_j\)인 \(dp[j]\)의 최댓값을 찾는다. 이 값은 구간 트리를 이용하면 빠르게 계산할 수 있다.

각 \(a_i\)와 \(b_i\)의 값들을 가지고 구간 트리를 만든 뒤, \(dp[i]\)에 해당하는 값을 구간 트리의 구간 (index(\(a_i\)), index(\(b_i\))]에 쿼리를 통해 찾은 값에 \(h_i\)를 더해 계산하고, \(b_i\)에 해당하는 위치(index(\(b_i)\))에 저장해두면 된다.

※ 여기서 index(x)는 구간 트리내에서 x값에 해당하는 원소의 번호이다. 

소스 코드

#include <bits/stdc++.h> using namespace std; typedef long long lint; const int MAX_N = 1e5; int n; int initSize; vector<int> init; vector<lint> tree; map<int, pair<int, lint>> m; // (b, (a, h)) int getIndex(int k) { return lower_bound(init.begin(), init.end(), k) - init.begin(); } void buildSegmentTree() { for (auto& it : m) { init.push_back(it.first); init.push_back(it.second.first); } sort(init.begin(), init.end()); init.erase(unique(init.begin(), init.end()), init.end()); initSize = init.size(); tree.assign(initSize * 4, 0); } lint query(int p, int left, int right, int qLeft, int qRight) { if (qLeft > right || qRight < left) return 0; if (qLeft <= left && right <= qRight) return tree[p]; int mid = (left + right) / 2; return max(query(p * 2, left, mid, qLeft, qRight), query(p * 2 + 1, mid + 1, right, qLeft, qRight)); } lint update(int p, int left, int right, int you, lint newValue) { if (you < left || you > right) return tree[p]; if (left == right) return tree[p] = newValue; int mid = (left + right) / 2; return tree[p] = max(update(p * 2, left, mid, you, newValue), update(p * 2 + 1, mid + 1, right, you, newValue)); } lint maxHeights() { lint ret = 0; for (auto& it : m) { int a = it.second.first; int b = it.first; lint h = it.second.second; int qRight = getIndex(b); int qLeft = min(qRight, getIndex(a) + 1); lint qValue = query(1, 0, initSize - 1, qLeft, qRight); lint newValue = qValue + h; update(1, 0, initSize - 1, qRight, newValue); ret = max(ret, newValue); } return ret; } int main() { scanf("%d", &n); for (int i = 0; i < n; i++) { int a, b; lint h; scanf("%d%d%lld", &a, &b, &h); auto where = m.find(b); if (where != m.end()) { if (where->second.first > a) { where->second.first = a; } where->second.second += h; } else { m[b] = { a, h }; } } buildSegmentTree(); printf("%lld", maxHeights()); }