[CodeForces - 1374E2] Reading Books (hard version)

난이도: 2500

 

태그

• Greedy
• Binary Search on Segment Tree

 

풀이

1. 누가 좋아하는가에 맞춰서 책을 나눠보자.

Alice가 좋아하는 책에는 1로, Bob이 좋아하는 책은 2로 분류해봅시다.

만약 둘 다 좋아한다면 1+2 = 3으로 분류하고, 둘 다 싫어한다면 0으로 분류하겠습니다.

 

2. 만약 1번 분류의 책을 2번 분류의 책보다 더 적게 읽을 거라면, 어떻게 책을 읽을 수 있을까?

\( C_i \) = \( i \)번 분류의 책을 "최소" 몇 권 읽을 것인지의 조건이라고 해봅시다.

그럼, [2]번 질문은 \( C_1 \le C_2 \)라는 추가 조건을 의미합니다.

 

여기에, \( C_1 \)을 고정시켜봅시다.

Alice는 적어도 \( k \)권의 재밌는 책을 읽어야 하므로, 3번 분류의 책은 \( \max(k - C_1, 0) \)권 이상 읽어야 합니다.

2번 분류의 책은 우리가 추가한 조건에 따라 \( C_1 \)권 이상 읽어야 하고요.

0번 분류의 책은 \( 0 \)권 이상 읽을 수 있고, 1번 분류의 책은 \( C_1 \)이 고정되었으므로 더 읽을 수 없습니다.

 

동일 분류의 책에서는 자명히 더 짧은 시간이 걸리는 책부터 읽어야 합니다.

그러니, 각각의 분류에서 반드시 읽어야 하는 책들은 가장 짧은 시간이 걸리는 책들로 채워봅시다.

 

그 뒤로는, 0, 2, 3번 분류에서 마음대로 책을 더 읽을 수 있습니다.

다만, 아직 정확히 \( m \)권의 책을 읽어야 한다는, 아직 하나의 조건이 남아있죠.

 

우리가 지금까지 읽은 책은 \( C_1 + C_1 + \max(k - C_1, 0) \)권입니다. 이를 토대로 남은 책의 수를 구할 수 있겠죠.

이 남은 책들 역시 자명히 시간이 짧게 걸리는 것부터 사용하면 되는데...

이제는 분류와 상관없이 빠르게 읽을 수 있는 \( k \)권의 시간 합을 계산해야 합니다.

 

그래서, 세그먼트 트리를 사용해볼 것입니다.

\( i \)번째 리프 정점에는 그 책을 아직 읽을 수 있는지에 따라 아래 두 경우 중 하나로 놓아봅시다.

// 읽을 수 없는 경우는 아래 두 경우 중 하나입니다.

// 1. 이미 읽은 책인 경우 / 2. 1번 분류에 속해서, 읽으면 안 되는 책인 경우

• 아직 그 책을 읽을 수 있다면: (읽는 시간이 \( i \)번째로 빠른 책의 시간, 1)
  // \( i \)번째 리프 정점에는 읽는 시간이 \( i \)번째로 빠른 책을 넣고 있습니다!
  // 이렇게 하면, 우리가 추가로 읽을 책들이 Prefix에 모이기 때문입니다.
• 그 책을 읽을 수 없다면: (0, 0)

그럼, 중간 정점에는 (그 구간 내의 책을 다 읽는 데 걸리는 시간, 그 구간 내의 책의 수)를 계산할 수 있겠죠.

 

이제 세그트리에서 이진 탐색을 해봅시다.

그 구간 내의 책의 수가 저장되어있으므로, 세그트리의 루트에서부터 \( x \)번째 책을 이진 탐색을 하며 내려갈 수 있습니다.

게다가, 그 구간 내의 책을 읽는 데 걸리는 시간이 저장되어있으므로, \( x \)권의 책을 읽는 시간을 계산하면서 내려갈 수 있습니다.

 

\( C_1 \)을 0권에서부터 1씩 늘려가면서 계산하면

반드시 읽어야 하는 책과 추가로 읽을 수 있는 책이 각 분류에서 최대 1개씩 변화되므로

스위핑 느낌나게 처리해주면 됩니다.

 

사실 위 과정만으로도 전체 문제를 풀 수 있습니다.

\( C_1 \ge C_2 \)인 경우는 Alice와 Bob을 반대로 처리해주면 되기 때문이죠.

 

3. 역추적을 위해 추가적으로 구현해야 할 사항들

우선, 계산 과정에서 답을 업데이트할 때마다 자기만 좋아하는 분류의 책을 더 적게 읽은 사람이 누구인지와 그게 몇 권이었는지 같이 업데이트하면,

 

나중에 모든 계산이 끝난 뒤 다시 (사람, 책의 수)의 상태로 되돌린 뒤, 그 과정을 다시 따라가보면서 선택된 책들의 번호를 하나씩 출력하면 됩니다.

 

4. 코드

int n, m, k;
vector<pi3> v[4];
 
const int N = 262144;
pi2 seg[524290];
void upd(int pos, pi2 val){ pos += N-1;
    //cout << "UPD " << pos-N+1 << " -> " << val.fr << ' ' << val.sc << endl;
    seg[pos] = val; pos >>= 1;
    while (pos){
        seg[pos].fr = seg[pos<<1].fr + seg[pos<<1|1].fr;
        seg[pos].sc = seg[pos<<1].sc + seg[pos<<1|1].sc;
        pos >>= 1;
    }
}
int kth(int k){
    int ni = 1;
    if (k == 0){ return 0; }
    int res = 0;
    while (ni < N){
        if (k <= seg[ni<<1].sc){ ni = ni<<1; }
        else{ k -= seg[ni<<1].sc; res += seg[ni<<1].fr; ni = ni<<1|1; }
        //cout << "KTH " << k << ' ' << res << ' ' << ni << endl;
    }
    return res + seg[ni].fr;
}
 
const int INF = 2147483647;
int ans = INF; pi2 anp = {0, 0};
void solve(int arg){
    for (int i = 1; i < N+N; i++){ seg[i] = {0, 0}; }
    for (pi3 p : v[0]){ upd(p.fr.sc, {p.fr.fr, 1}); }
    for (pi3 p : v[2]){ upd(p.fr.sc, {p.fr.fr, 1}); }
    for (pi3 p : v[3]){ upd(p.fr.sc, {p.fr.fr, 1}); }
    int l0 = v[0].size(), l1 = v[1].size(), l2 = v[2].size(), l3 = v[3].size();
    int res = 0;
    int c3 = min(k, l3);
    for (int i = 0; i < c3; i++){
        pi2 p = v[3][i].fr;
        upd(p.sc, {0, 0}); res += p.fr;
    }
    for (int c1 = 0; c1 <= m; c1++){
        if (c1 > 0){
            if (c1 > l1 || c1 > l2){ break; }
            pi2 p1 = v[1][c1-1].fr, p2 = v[2][c1-1].fr;
            upd(p1.sc, {0, 0}); res += p1.fr;
            upd(p2.sc, {0, 0}); res += p2.fr;
        }
        if (c3 > max(k-c1, 0)){
            pi2 p3 = v[3][c3-1].fr;
            upd(p3.sc, {p3.fr, 1}); res -= p3.fr; c3 -= 1;
        }
        if (c3 < max(k-c1, 0)){ continue; }
        int cnt = c1+c1+c3;
        if (cnt > m){ break; } if (n-(l1-c1) < m){ continue; }
        int val = res + kth(m-cnt);
        //cout << arg << ' ' << c1 << ' ' << c3 << ' ' << cnt << " -> " << res << ' ' << val << endl;
        if (val < ans){ ans = val; anp = {arg, c1}; }
    }
}
 
void Main(){
    cin >> n >> m >> k;
    vector<pi2> srt;
    for (int i = 1; i <= n; i++){
        int x, a, b; cin >> x >> a >> b;
        int idx = 2*a + b;
        v[idx].push_back({{x, 0}, i});
        srt.push_back({x, idx});
    }
    for (int i = 0; i < 4; i++){ sort(all(v[i])); } sort(all(srt));
    vector<pi3>::iterator it[4]; for (int i = 0; i < 4; i++){ it[i] = v[i].begin(); }
    int i = 0; for (pi2 p : srt){ i += 1;
        (*it[p.sc]).fr.sc = i; it[p.sc]++;
    }
    solve(0); swap(v[1], v[2]); solve(1);
    if (ans == INF){ cout << -1; return; } else{ cout << ans << endl; }
    if (anp.fr == 0){ swap(v[1], v[2]); }
    int cnt = anp.sc;
    int c0 = 0, c1 = cnt, c2 = cnt, c3 = max(k-cnt, 0);
    for (int i = 0; i < c1; i++){ cout << v[1][i].sc << ' '; }
    for (int i = 0; i < c2; i++){ cout << v[2][i].sc << ' '; }
    for (int i = 0; i < c3; i++){ cout << v[3][i].sc << ' '; }
    vector<pi2> val;
    for (int i = c0; i < v[0].size(); i++){ val.push_back({v[0][i].fr.fr, v[0][i].sc}); }
    for (int i = c1; i < v[1].size(); i++){ val.push_back({v[1][i].fr.fr, v[1][i].sc}); }
    for (int i = c2; i < v[2].size(); i++){ val.push_back({v[2][i].fr.fr, v[2][i].sc}); }
    for (int i = c3; i < v[3].size(); i++){ val.push_back({v[3][i].fr.fr, v[3][i].sc}); }
    sort(all(val));
    int lft = m - (c1+c2+c3);
    for (int i = 0; i < lft; i++){ cout << val[i].sc << ' '; }
}