[BOJ] 19646 - Random Generator

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

BOJ - Random Generator

문제 유형: 세그먼트 트리, 이분 탐색

문제 난이도: Platinum V

시간 제한: 1초

메모리 제한: 1024MB

 

문제

국렬이는 1부터 N까지의 양의 정수로 이루어진 순열을 주어진 양의 정수 w1 ... , wN를 이용해서 무작위로 만들 것이다. 다음은 무작위로 순열을 만드는 방법이다.

1. 1부터 N까지의 양의 정수 i (1 ≤ i ≤ N)를 연속적으로 wi개씩 배치한다.
2. 현재 배치된 양의 정수의 총 개수를 W라고 하자. 1부터 W까지의 양의 정수들 중에서 균등하게 숫자 하나 pi를 선택한다.
3. pi번째 수를 순열에 추가한다.
4. 순열에 추가한 수들을 전부 지우고, 남은 수가 없을 때까지 2부터 4의 과정을 거친다.

w1부터 wN까지, p1부터 pN까지의 수가 주어졌을 때, 나오는 순열을 구해 보자.

 

입력

첫 번째 줄에는 N (1 ≤ N ≤ 200,000)이 주어진다.

두 번째 줄에는 1,000 이하의 양의 정수 w1부터 wN까지가 주어진다.

세 번째 줄에는 양의 정수 p1부터 pN까지가 주어진다. p1부터 pN까지로 수열을 만들 수 있는 경우만 주어진다.

 

출력

최종적으로 나오는 1부터 N까지 양의 정수들로 이루어진 수열을 구하자.

 

풀이

세그먼트 트리를 이용해서 구간의 합을 최적화했다.

세그먼트 트리로 구간합을 저장하도록 하였다.

findKth()로 k번째 수가 어느 숫자인지 이분 탐색으로 구하고,

update()함수를 통해서 구간합을 비워준다.

 

각 p값에 대해서 findKth -> 출력 -> update 사이클로 구현했다.

 

세부 구현을 아래를 확인.

 

코드

C++

#include <bits/stdc++.h>
#define MAX 200001

using namespace std;


int w[MAX];
int seg[MAX*4];
int n;

void build(int idx, int l, int r) {
    if(l == r) {
        seg[idx] = w[l];
        return;
    }

    int mid = (l + r) >> 1;
    build(idx *2, l, mid);
    build(idx *2 + 1, mid+1, r);
    seg[idx] = seg[idx * 2] + seg[idx * 2 + 1];
}

void update(int idx, int target, int l, int r) {
    if(l == r) {
        w[l] = 0;;
        seg[idx] = 0;
        return;
    }

    int mid = (l + r) >> 1;
    if(target <= mid) {
        update(idx *2, target, l, mid);
    } else {
        update(idx*2+1, target, mid+1, r);
    }

    seg[idx] = seg[idx * 2] + seg[idx*2 + 1];
}

int findKth(int idx, int l, int r, int k) {
    if (l == r) return l;
    int leftSum = seg[idx*2];
    int mid = (l + r) >> 1;
    if(k <= leftSum) return findKth(idx*2, l, mid, k);
    else return findKth(idx*2+1, mid+1, r, k-leftSum);
}

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

    cin >> n;

    for(int i = 1; i <= n; i++) {
        cin >> w[i];
    }
    build(1, 1, n);

    for(int i = 0; i < n ; i++) {
        int p;
        cin >> p;

        int idx = findKth(1, 1, n, p);
        cout << idx << " ";
        update(1, idx, 1, n);
    }
    cout << "\n";

    return 0;
}