[BOJ] 27232 - 청소

https://www.acmicpc.net/problem/27232

BOJ - 청소

문제 유형: 집합, 슬라이딩 윈도우

문제 난이도: Gold I

시간 제한: 2초

메모리 제한: 1024MB

 

문제

준석이는 청소 업체에 다니고 있다. 준석이가 청소할 장소는 1번부터 N번까지 차례로 번호가 붙은 일렬의 1N개의 구역으로 나누어져 있다. i번 구역은 i-1번과 i+1번 구역과 인접해 있어, 두 인접한 구역 사이를 이동하려면 1만큼 걸어야 한다.

각 구역에는 우선순위가 있다. i번 구역의 우선순위 A_i는 1이상 N이하의 정수로 이 값이 클수록 우선순위가 높다. 임의의 두 구역의 우선순위는 항상 다르다.

준석이는 오늘 이 구역 중 K개의 구역을 먼저 청소하려고 한다. 단, 준석이가 청소하는 K개의 구역은 반드시 연속해야 한다. 또한 청소할 때 선택한 구역들 내에서는 우선순위가 높은 구역부터 낮은 구역 순서대로 이동하며 청소해야 한다.

두 구역이 멀리 떨어져 있으면 이동하는 시간이 오래 걸리기 때문에, 준석이는 이동 거리의 합이 최소화되도록 연속한 K개의 구역을 선택하려고 한다. 이동 거리가 최소가 되도록 구역을 선택했을 때 총 이동 거리를 출력한다.

 

입력

첫째 줄에 청소할 구역의 길이 N과 오늘 청소할 구역의 개수 K가 공백으로 구분되어 주어진다.

둘째 줄에 각 구역의 우선순위 A_1,A_2,... ,A_N이 공백으로 구분되어 주어진다.

 

출력

 K개의 구역을 선택했을 때 가능한 총 이동 거리 중 최솟값을 출력한다.

 

풀이

k개의 구역을 선택을 슬라이딩 윈도우를 통해서 구하면 된다.

그러나, 문제는, 매번 전체 구역의 계산하면, O(N*K)가 걸리고, 이는 TLE이다.

 

매번의 윈도우마다, O(log K)로 구해내야 하는데, 어떻게 할 수 있을까?

정답은 set을 이용하는 것이다.

set내부에는 레드블랙트리로 구현되어있고, 이전 값과 다음 값을 O(1)에 구할 수 있다.

 

즉, 매번 전체의 이동거리를 구할 필요 없이, 변경사항에 대해서만 이동거리를 업데이트 시킬 수 있다.

세부구현은 코드 참조..

 

코드

C++

#include <bits/stdc++.h>

using namespace std;
using ll = long long;

int main(int argc, char* argv[]) {
    ios::sync_with_stdio(0);
    cin.tie(0);

    int n, k;
    cin >> n >> k;

    vector<int> a(n + 1);
    vector<int> pos(n + 1);
    for (int i = 1; i <= n; ++i) {
        cin >> a[i];
        pos[a[i]] = i;
    }

    set<int> s;
    ll ans = 1e18;
    ll dist = 0;

    auto add = [&](int x) {
        auto it = s.insert(x).first;
        if (it != s.begin()) {
            dist += llabs((ll)pos[*it] - pos[*prev(it)]);
        }
        auto itn = next(it);
        if (itn != s.end()) {
            dist += llabs((ll)pos[*it] - pos[*itn]);
        }
        if (it != s.begin() && itn != s.end()) {
            dist -= llabs((ll)pos[*itn] - pos[*prev(it)]);
        }
    };

    auto remove = [&](int x) {
        auto it = s.find(x);
        auto itn = next(it);
        if (it != s.begin()) {
            dist -= llabs((ll)pos[*it] - pos[*prev(it)]);
        }
        if (itn != s.end()) {
            dist -= llabs((ll)pos[*it] - pos[*itn]);
        }
        if (it != s.begin() && itn != s.end()) {
            dist += llabs((ll)pos[*itn] - pos[*prev(it)]);
        }
        s.erase(it);
    };

    for (int i = 1; i <= n; ++i) {
        add(a[i]);
        if (s.size() > k) remove(a[i - k]);
        if (s.size() == k) ans = min(ans, dist);
    }
    cout << ans << "\n";

    return 0;
}