백준 10076번 휴가

문제 요약

크기 n인 배열이 주어진다. 그리고 시작위치 start와 휴가일수 d가 주어진다. 배열의 값 $a_i$는 i번째 지점에 위치한 관광지 수이다.

배열을 한 칸 이동하는데에 1일이 사용되고 현재 위치한 곳의 관광지를 전부 보는 데에 1일이 사용된다. 한 번 본 곳은 다시 봐도 의미가 없다.

이 때, 돌아볼 수 있는 최대 관광지 수를 구해야 한다.

 

n은 최대 10만이며 d는 최대 2n + n//2이다.

풀이

제일 먼저 할 수 있는 관찰은 왔다갔다를 여러번 하는 것은 절대 최적일 수가 없다.

왼쪽으로 갔다가 오른쪽으로 갔다가 다시 왼쪽으로 가는 것은 휴가일수를 낭비하기만 한다.

따라서, 한 방향으로 쭉 갔다가 한번 꺾어서 쭉 가는 것만 고려한다고 하자.

 

그러면 왼쪽으로 갔다가 오른쪽으로 가는 거랑 오른쪽으로 갔다가 왼쪽으로 가는거랑 두 가지가 있는데 다른 한가지를 해결하면 배열을 뒤집어서 진행하면 되기 때문에 왼쪽으로 갔다가 오른쪽으로 가는 것을 기준으로 생각하자.

 

왼쪽으로 갔다가 오른쪽으로 가는데 왼쪽으로 가다가 오른쪽으로 꺾는 지점을 $l$, 마지막에 위치할 오른쪽 지점을 $r$이라고 하자.

$l \le start, r \ge start$로 생각하도록 하자. 그러면 $[l,r]$을 움직인다고 할 때, 움직이는 거리는 $(start-l)+(r-l)$이 된다.

이제 휴가일 수 $d$가 주어졌을 때, $[l,r]$에서 돌아볼 수 있는 지점의 수는 $d - (start-l)-(r-l)$이 됐다. $d-start-r+2l$이라고 하자.

당연히 이 값이 0보다 작으면 점수를 얻을 수 없다.

 

이제 $[l,r]$이 고정됐을 때 볼 수 있는 지점의 수가 바로 결정됨을 알았다. 그 수를 $k$라고 하자.

그러면 하나의 구간 $[l,r]$에서 돌아볼 수 있는 최대 관광지 수는 $a_l, a_{l+1}, \cdots , a_r$중에서 가장 큰 $k$개의 합이다. 이것은 세그먼트 트리로 구할 수 있다.

여기까지의 풀이는 $O(N^2logN)$이다. 봐야될 구간이 $O(N^2)$이기 때문이다.

 

특정 $l$에 대해서 최적의 값을 갖게되는 $r$을 $opt_l$이라고 하자.

$l<l'$일 때, $opt_l \le opt_{l'}$임이 꽤 간단하게 나온다. 간단한 직관만 설명하면 방문 구간이 왼쪽에서 줄었으면 오른쪽에서도 줄어들 일은 절대 없다는 것이다.

어쨌든 이 사실을 이용하면 분할정복 최적화를 사용할 수 있다. 따라서, 총 시간복잡도가 $O(Nlog^2N)$으로 줄어들었고 문제를 해결할 수 있다.

 

가장 큰 $k$개의 합을 구하는 데에 나는 PST를 사용헀다. 메모리 초과에는 적절히 주의해야한다. PST말고도 쌩 세그먼트 트리를 사용해도 된다. 이 경우에는 업데이트 처리가 귀찮을 거 같긴 하다.

 

합을 구하는 방법은 주어진 배열을 내림차순으로 정렬하고 인덱스 1부터 N까지 해당 위치의 원소에 1씩 키운다. 이 다음은 K번째 수를 구하는 방법과 동일하다.

#include<bits/stdc++.h>
using namespace std;
using ll = long long;
const int MXN = 1e5;
struct Node {
    int l, r; 
    ll sum;
    int cnt;
    Node() : cnt(0), sum(0) {};
};
vector<ll> v;
vector<int> a;
vector<int> roots(MXN+1);
int N, S, D;
vector<Node> tree(MXN*20);
ll ans;
int node_cnt;
void init() {
    node_cnt = v.size() * 2;
    for(int i=1;i<v.size();++i) {
        tree[i].cnt = tree[i].sum = 0;
        tree[i].l = i << 1;
        tree[i].r = i << 1 | 1;
    }
}

void update(int node, int s, int e, int idx, ll val) {
    tree[node].cnt += val;
    tree[node].sum += val * v[idx];
    if(s != e) {
        int mid = s + e >> 1;
        int n1 = tree[node].l, n2 = tree[node].r;
        if(idx <= mid) {
            tree[node_cnt] = tree[n1];
            tree[node].l = node_cnt++;
            update(tree[node].l, s, mid, idx, val);
        }
        else {
            tree[node_cnt] = tree[n2];
            tree[node].r = node_cnt++;
            update(tree[node].r, mid+1, e, idx, val);
        }
    }
}

ll query(int s, int e, int k) {
    int l = 0, r = v.size() - 1;
    ll ret = 0;
    while(l != r) {
        int mid = l + r >> 1;
        int lsz = tree[tree[e].l].cnt - tree[tree[s].l].cnt;
        ll lsum = tree[tree[e].l].sum - tree[tree[s].l].sum;
        if(lsz >= k) {
            s = tree[s].l; e = tree[e].l;
            r = mid;
        }
        else {
            ret += lsum;
            k -= lsz;
            s = tree[s].r; e = tree[e].r;
            l = mid + 1;
        }
    }
    assert(l == r);
    ret += v[l] * min(k, (int)tree[e].cnt - (int)tree[s].cnt);
    return ret;
}

void make_tree() {
    roots[0] = 1;
    for(int i=1;i<=N;++i) {
        roots[i] = node_cnt;
        tree[node_cnt++] = tree[roots[i-1]];
        update(roots[i], 0, v.size()-1, a[i], 1);
    }
}

void dnc(int s, int e, int opts, int opte) {
    if(s > e) return ;
    ll ret = 0;
    int opt_mid = opts;
    int mid = s+e >> 1;
    for(int i=opts;i<=opte;++i) {
        int k = max(0, D - S - i + 2*mid);
        k = min(k, i - mid + 1);
        ll tmp = query(roots[mid-1], roots[i], k);
        if(tmp > ret) {
            ret = tmp; opt_mid = i;
        }
    }
    ans = max(ret, ans);
    dnc(s, mid-1, opts, opt_mid);
    dnc(mid+1, e, opt_mid, opte);
}

int main() {
    cin.tie(nullptr); ios::sync_with_stdio(false);
    cin >> N >> S >> D; ++S;
    a.resize(N+1);
    for(int i=1;i<=N;++i) {
        cin >> a[i]; v.push_back(a[i]);
    }
    sort(v.begin(), v.end(), greater<ll>());
    v.erase(unique(v.begin(), v.end()), v.end());
    for(int i=1;i<=N;++i) a[i] = lower_bound(v.begin(), v.end(), a[i], greater<ll>()) - v.begin();
    init();
    make_tree();
    dnc(1, S, S, N);
    reverse(a.begin()+1, a.begin()+N+1); S = N - S + 1;
    init();
    make_tree();
    dnc(1, S, S, N);
    cout << ans << '\n';
    return 0;
}