2016 SCPC 2차 온라인 예선 3번 - 땅 나누기

귀찮아서 안 쓰고 있었는데, O(nlogn)임에도 계속 시간초과를 받았던게 자꾸 생각나서 늦게나마 해법(결국 Fast I/O까지 써서 시간초과 X)을 올린다 -_-..

문제

철광석과 구리 광석이 묻혀있는 땅이 있다. 여기서 N(1 ≤ N ≤ 100,000)개의 광물(점으로 간주한다)의 좌표가 주어질 때, 어떤 점 P를 원점으로 잡고 1사분면의 구리 광석 개수, 2사분면의 철광석 개수, 3사분면의 구리 광석 개수, 그리고 4사분면의 철광석 개수를 셌을 때의 값을 최대화하도록 P를 정하고, 이때의 합을 출력하자.

파란색 점을 철광석, 갈색 점을 구리 광석이라고 생각하자.

P의 위치를 위와 같이 잡을 경우 합은 3(1사분면) + 3(2사분면) + 3(3사분면) + 2(4사분면) = 11이 된다.

*각 광물의 좌표는 모두 1 ≤ X,Y ≤ N 범위의 정수(자연수)이고, 모든 광물의 X좌표는 서로 다르고, Y좌표도 서로 다르다. 즉, 어떤 수평선 위에도 광물은 최대 한 개 존재하고, 어떤 수직선 위에도 광물은 최대 한 개 존재한다. 그리고, 점 P의 좌표는 실수를 가질 수 있다(즉, 경계선에 놓이는 광물은 없다고 보아도 된다).

풀이

우선 이 풀이를 이해하기 위해선 <Segment Tree> 라는 자료구조에 대한 이해가 필요하다 [링크1] [링크2].

Segment Tree를 처음 배울 때 보통 구간의 합을 구하는 예제를 통해 배우게 되는데, 이를 조금 확장하면 어떤 배열의 Prefix sum, 즉 A[1]+A[2]+...+A[i] (1 ≤ i ≤ N) 의 합이 최대가 되는 i와 그때의 합을 찾을 수 있다. 리프 노드의 최대 접두사 합은 자기 자신의 값으로 하고, 리프가 아닌 노드의 최대 접두사 합은 "왼쪽 자식의 최대 접두사 합"과 "왼쪽 자식의 합 + 오른쪽 자식의 최대 접두사 합" 중 큰 것임이 자명하기 때문에 접두사 합을 저장하는 배열을 하나 더 만들거나, 구조체를 이용해서 세그먼트 트리를 만들면 된다.

#include<cstdio>
#include<algorithm>
using namespace std;
const int MAXN(131072);

struct PrefixSumTree {
    struct Node {
        int sum, max_prefix;
        inline Node(): sum(0), max_prefix(0) {}
        inline Node(int a): sum(a), max_prefix(a) {}
        inline void Combine(const Node& le, const Node& ri) {
            sum = le.sum + ri.sum;
            max_prefix = max(le.max_prefix, le.sum+ri.max_prefix);
        }
    };

    int n;
    Node *nodes;

    inline PrefixSumTree(int *arr, int n): n(n) {
        nodes = new Node[n << 1];

        for (int i = 0; i < n; i++) {
            nodes[i+n] = Node(arr[i]);
        }

        for (int i = n - 1; i > 0; --i) {
            nodes[i].Combine(nodes[i<<1], nodes[i<<1|1]);
        }
    }

    inline void Update(int p, int value) {
        nodes[p+n] = Node(value);

        for (p += n; p > 1; p >>= 1) {
            int parent = p >> 1, left = p, right = p;
            if (p&1) left ^= 1;
            else right ^= 1;
            nodes[parent].Combine(nodes[left], nodes[right]);
        }
    }

    inline int QueryMaximum() {
        return nodes[1].max_prefix;
    }

    inline ~PrefixSumTree() {
        delete[] nodes;
    }
};

이제 위의 자료구조를 이용해서 문제를 해결하는 방법에 대해 살펴보자.

점 P의 X 좌표가 특정한 값(이 값을 Px라고 하자)로 고정되어 있을 경우, 최적의 Y 좌표를 쉽게 구하는 방법은 무엇일까?

직관적으로 쉽게 알 수 있는데, 아래 그림과 같이 배열을 하나 만들고, X좌표가 Px보다 작은 철광석이 있는 Y좌표에는 -1, 구리 광석이 있는 Y좌표에는 +1, 그리고 X좌표가 Px보다 크거나 같은 철광석이 있는 Y좌표에는 +1, 구리 광석이 있는 좌표에는 -1을 준 다음, 접두사 합이 최대가 되는 지점이 바로 Px의 Y좌표가 된다(증명은 귀류법으로 간단히 할 수 있다).

맨 아래에서부터 (-1) + 1 + (-1) + (-1) + 1 + 1 + 1 = 1

이를 처음에 설명한 트리를 응용하여 모든 X좌표(N개(정확힌 N+1개)의 X좌표)에 대해 빠르게(O(logN)) 최적의 Y좌표를 구할 경우, 모든 P의 후보 지점을 살펴보는 데에 O(NlogN)의 시간이 걸리므로 최대 100,000개의 점을 처리하기엔 부족함이 없다.

우선 Px를 나타내는 가상의 수직선을 상상해보고, 다음 그림들을 보자. 

 -> 

이런 상황으로 초기화를 한다(px = 0) - O(NlogN) -> 여기서 최적의 y를 찾는다 - O(1) -> 합을 구한다 - O(1).
(합은 생각해보면, 간단하게 (최대 접두사 합) + (P 오른쪽의 구리 개수) + (P 왼쪽의 철 개수)이다)

이제 수직선을 한 칸 이동시켜서(Px = 1) 반복해보자.

수직선을 이동할 때 갱신해야 하는 것은 세 가지이다:

1. 세그먼트 트리에서 해당 X좌표에 있는 점의 부호를 바꿔준다 - O(logN).

2. 해당 점이 철광석일 경우 P 왼쪽에 있는 철 개수를 저장하는 변수(LEFT_IRON) - O(1)

3. 해당 점이 구리 광석일 경우 P 오른쪽에 있는 구리 개수를 저장하는 변수(RIGHT_COPPER) - O(1).

다시 같은 방식으로 Px=2, ...., Px = N까지 모두 확인해가며 합의 최댓값을 갱신해주면 정답을 구할 수 있다.

#include<cstdio>
#include<algorithm>
using namespace std;
const int MAXN(131072);

struct PrefixSumTree {
    struct Node {
        int sum, max_prefix;
        inline Node(): sum(0), max_prefix(0) {}
        inline Node(int a): sum(a), max_prefix(a) {}
        inline void Combine(const Node& le, const Node& ri) {
            sum = le.sum + ri.sum;
            max_prefix = max(le.max_prefix, le.sum+ri.max_prefix);
        }
    };

    int n;
    Node *nodes;

    inline PrefixSumTree(int *arr, int n): n(n) {
        nodes = new Node[n << 1];

        for (int i = 0; i < n; i++) {
            nodes[i+n] = Node(arr[i]);
        }

        for (int i = n - 1; i > 0; --i) {
            nodes[i].Combine(nodes[i<<1], nodes[i<<1|1]);
        }
    }

    inline void Update(int p, int value) {
        nodes[p+n] = Node(value);

        for (p += n; p > 1; p >>= 1) {
            int parent = p >> 1, left = p, right = p;
            if (p&1) left ^= 1;
            else right ^= 1;
            nodes[parent].Combine(nodes[left], nodes[right]);
        }
    }

    inline int QueryMaximum() {
        return nodes[1].max_prefix;
    }

    inline ~PrefixSumTree() {
        delete[] nodes;
    }
};

struct Ore {
    int x, y, z;
    inline Ore () {}
    inline Ore(int x, int y, int z): x(x), y(y), z(z) {}
    inline bool operator < (const Ore& j) const {
        return x < j.x;
    }
};

int n;
Ore total[MAXN];
PrefixSumTree *prefixsum_tree;

inline int Process(int base)
{
    int ans = 0;
    for (int i = 0; i <= n; i++) {
        if (i > 0) {
            int j = i-1;
            if (total[j].z) {
                prefixsum_tree->Update(total[j].y, -1);
                ++base;
            } else {
                --base;
                prefixsum_tree->Update(total[j].y, 1);
            }
        }

        int sum = base + prefixsum_tree->QueryMaximum();
        ans = max(ans, sum);
    }

    return ans;
}

inline void readn(register int *n)
{
    register char c;
    (*n) = 0;

    while (~(c = getc_unlocked(stdin))) {
        if ('0' <= c && c <= '9') break;
    }

    do {
        if (c < '0' || c > '9') break;
        (*n)*=10, (*n)+=c&~48;
    } while (~(c = getc_unlocked(stdin)));
}

int main()
{
    int T, TC = 1;
    setbuf(stdout, NULL);

    for (readn(&T); T--; TC++) {

        int arr[MAXN] = {}, base = 0;
        readn(&n);

        printf("Case #%d\n", TC);

        for (int i = 0; i < n; i++) {
            readn(&total[i].x);
            readn(&total[i].y);
            readn(&total[i].z);
            if (total[i].z) {
                arr[total[i].y] += 1;
            } else {
                arr[total[i].y] -= 1;
                ++base;
            }
        }

        sort(total, total+n);

        prefixsum_tree = new PrefixSumTree(arr, MAXN);
        printf("%d\n", Process(base));

        delete prefixsum_tree;
    }
    return 0;
} /* syntax highlighter: http://markup.su/ */

*10.23: 정식 채점이 가능해짐에 따라 코드 오탈자 수정(정식 채점: http://codeground.org)

테스트 데이터

랜덤하게 만들어본 테스트 데이터와 이에 해당하는 정답 데이터로, 이 데이터에 대한 정답이 0.5초안에 모두 나오지 않는다면 일단 시간초과는 확실한걸로..

sample.in.lzma

sample.out.lzma

(압축 해제를 해야 합니다..)