D. Learning to Paint (CodeTON Round 8 Div. 1 + Div. 2)

• 링크

https://codeforces.com/contest/1942/problem/D

• 내가 푼 방법

거의 다 접근했는데, 정말 아쉽다.

이 문제의 핵심은 현재 인덱스에서 범위를 왼쪽으로 쭉 늘려가면서 최적의 값을 구하는 것이다. 사실 여기서 끝나면 되게 쉬운 문제이지만, 가장 큰 k개를 뽑아내야 하므로, 현재 인덱스에서 얻어낼 수 있는 최댓값을 최대 k개 구해놓아야 한다.

그러므로, 모든 index에는 최대 k개의 숫자가 있을 수 있다. 그렇다면 딱 봐도 n n k 번을 돌아야 할 것 같다. 이렇게 되면 무조건 터진다.

그러니, 꼼수를 써주어야 한다. 우선순위 큐를 사용하는 것이다. 여기까지 접근했는데... 한 발짝을 더 못 갔다.

이전 인덱스에 저장된 리스트들도 전부 다 내림차순으로 정렬되어 있다. 정확히는 증가하지 않는 순으로 말이다.

그러므로 일단, 현재 인덱스 바로 이전부터 0번째까지 쭈욱 가면서 값을 구해준다. 이때 쭈욱 가는 인덱스는 j, 그리고 현재의 인덱스는 i라고 했을 때, dp[j][0]+arr[j+2][i]가 일단은 최적의 값이다.

하지만, dp[j]에는 하나의 수만 존재하는 것이 아닐 수도 있다. 리스트이기 때문이다.

그래서 현재 dp[j]의 몇 번째 인덱스의 값을 넣었는지와 내가 넣은 값의 인덱스를 알기 위해 j도 같이 넣어준다.

그러므로 q.push({dp[j][0]+arr[j+2][i],j,0}) 라는 수식이 나올 수 있는 것이다.

이렇게 하면 끝났다. 그냥 dp[i] 의 사이즈가 k가 되거나 혹은 q의 사이즈가 0이 되면 끝내면 된다.

또한, 내부에서 계속해서 큰 값을 뽑아내고, dp[j] 의 다음 인덱스를 넣는다. 즉 dp[j][0] 을 처음에 넣었으니, 다음에는 dp[j][1] 을 넣는 것이다.

또한 여기서 중요한 점은 j < 0인 값은 continue 해주는 것이다. 이유는 간단하다 인덱스를 벗어나기 때문이다. ㅋㅋㅋ c++에서 Out Of Index 를 잘 잡아주지 않아서 그런가? 이것 때문에 꽤 해맸다.

이렇게 한 뒤에 dp[n] 에 들어있는 값들 모두를 구하면 끝까지 진행하면서 큰 순서대로 k 개 의 수들을 출력할 수 있다.

• 시간 복잡도

O((n + k) * log n)

• Correct Code

#include <algorithm>
#include <cmath>
#include <cstdio>
#include <iostream>
#include <vector>
#include <queue>

using namespace std;

struct comp {
    bool operator()(vector<int>& o1, vector<int>& o2) {
        return o1[0] < o2[0];
    }
};

void solve() {
    int n, k;
    cin >> n >> k;
    vector<vector<int>> arr(n + 1, vector<int>(n + 1, 0));
    for (int i = 1; i <= n; i++) {
        for (int j = i; j <= n; j++) {
            cin >> arr[i][j];
        }
    }
    vector<vector<int>> dp(n + 1);
    dp[0].push_back(0);
    for (int i = 1; i <= n; i++) {
        priority_queue<vector<int>, vector<vector<int>>, comp> q; //{value,idx,dp idx}
        q.push({dp[i - 1][0], i - 1, 0});
        q.push({arr[1][i], -1, 0});
        for (int j = i - 2; 0 <= j; j--) {
            q.push({dp[j][0] + arr[j + 2][i], j, 0});
        }
        while (!q.empty() && dp[i].size() < k) {
            vector<int> vv = q.top();
            q.pop();
            int v = vv[0], iIdx = vv[1], jIdx = vv[2];
            dp[i].push_back(v);
            if (iIdx < 0 || jIdx + 1 == dp[iIdx].size()) {
                continue;
            }
            if (iIdx == i - 1) {
                q.push({dp[i - 1][jIdx + 1], iIdx, jIdx + 1});
            } else {
                q.push({dp[iIdx][jIdx + 1] + arr[iIdx + 2][i], iIdx, jIdx + 1});
            }
        }
    }
    for (auto v : dp[n]) {
        cout << v << " ";
    }
    cout << "\n";
}

int main(void) {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);
    freopen("d.input.txt", "r", stdin);
    int T;
    cin >> T;
    while (T-- > 0) {
        solve();
    }
    return 0;
}