백준 21399번: 장난감 묶음 할인 [C++]

작성일 | In
플래티넘 V 플래티넘 V

문제

장난감 묶음 할인

풀이

Albert가 팔기로 고른 장난감의 번호를 $s$라 합시다. 이때 Alice는 왼쪽 $[1, s - 1]$을 $(s - 1) / 3$-묶음 할인받고, 오른쪽 $[s + 1, n]$을 $(n - s)/3$-묶음 할인받는 것이 최선입니다.

왼쪽과 오른쪽에서 각각 가장 큰 원소를 뽑아내기 위한 우선순위 큐(최대 힙)와 할인된 원소들을 저장할 트리 집합을 하나씩 만듭니다. 이를 lheap, rheap, ldisc, rdisc라 부르겠습니다. 원소는 항상 (가격, 인덱스) 쌍으로 관리합니다.

초기에 $s = 0$이고 lheap과 ldisc는 비어있으며 rdisc에는 $[1, n]$에서 가장 큰 원소 $(n - 1)/3$개가, rheap에는 나머지가 들어있습니다.

$s$에 $3$을 더할 때 마다 아래를 반복하면서 가격을 구해주면 됩니다.

  1. $s - 3 \le i \le s$인 $i$가 rdisc에 존재하는 경우
    $i$는 더 이상 오른쪽에서 할인받을 수 없으므로 rdisc에서 제거한 뒤 다시 rheap에서 원소를 뽑아 할인받습니다.
    이때 rheap에서 뽑은 원소의 인덱스가 $s$보다 작거나 같지는 않은지 검사해야 합니다.
  2. $s - 3 \le i < s$인 $i$가 ldisc의 가장 작은 원소보다 가격이 큰 경우
    이때는 ldisc의 가장 작은 원소보다 $i$를 할인받는 것이 이득이므로 ldisc의 가장 작은 원소를 꺼내어 lheap에 넣은 뒤 $i$를 할인받습니다.
  3. $s - 3 \le i < s$인 $i$가 ldisc의 가장 작은 원소보다 가격이 작은 경우
    나중에 할인받을 수도 있으므로 lheap에 삽입합니다.
  4. 왼쪽편에 원소 3개가 추가되었으므로 lheap에서 하나를 꺼내어 할인받습니다.
  5. 오른쪽편에 원소 3개가 제거되었으므로 rdisc에서 가장 작은 원소를 제거합니다.

시간 복잡도는 $\mathcal{O}(Tn\log n)$입니다.

코드

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
#include <bits/stdc++.h>
#ifndef ONLINE_JUDGE
#define ASSERT(x) assert(x)
#else
#define ASSERT(ignore) ((void)0)
#endif

using namespace std;

using ll = long long;

ll price[30303];

void solve() {
    int n;
    cin >> n;

    for (int i = 0; i < n; i++)
        cin >> price[i];

    priority_queue<pair<ll, int>> lheap, rheap;
    set<pair<ll, int>> ldisc, rdisc;

    int select = 0;
    ll total = price[select];
    for (int i = 1; i < n; i++) {
        total += price[i];
        rheap.emplace(price[i], i);
    }

    for (int i = 0; i < (n - 1) / 3; i++) {
        auto [p, idx] = rheap.top();
        rheap.pop();
        total -= p;
        rdisc.emplace(p, idx);
    }

    pair<ll, int> answer = pair(total, select);
    // cout << total << '\n';

    for (select = 3; select < n; select += 3) {
        for (int i = select - 3; i <= select; i++) {
            if (rdisc.erase(pair(price[i], i))) {
                total += price[i];

                while (not rheap.empty()) {
                    auto [p, idx] = rheap.top();
                    rheap.pop();

                    if (idx <= select)
                        continue;

                    total -= p;
                    rdisc.emplace(p, idx);
                    break;
                }
            }

            if (i < select) {
                if (not ldisc.empty() and ldisc.begin()->first < price[i]) {
                    total += ldisc.begin()->first;
                    lheap.push(*ldisc.begin());
                    ldisc.erase(ldisc.begin());

                    total -= price[i];
                    ldisc.emplace(price[i], i);
                } else {
                    lheap.emplace(price[i], i);
                }
            }
        }

        // left side discount add
        total -= lheap.top().first;
        ldisc.emplace(lheap.top());
        lheap.pop();

        // right side discount removal
        if (not rdisc.empty()) {
            total += rdisc.begin()->first;
            rheap.push(*rdisc.begin());
            rdisc.erase(rdisc.begin());
        }

        // cout << total << '\n';
        if (answer.first < total)
            answer = pair(total, select);
    }

    cout << answer.second << ' ' << answer.first << '\n';
}

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

    int T;
    cin >> T;
    while (T--)
        solve();

    return 0;
}

백준 21399번: 장난감 묶음 할인 [C++]