문제 링크

세그트리 구현보단 그걸로 수열 뽑는 로직 발상이 좀 어렵다.

풀이

수열 내에서 특정 인덱스에 있는 값을 삭제하고 그 값의 배열 인덱스를 출력하는 문제다. 수열 요소를 삭제할 때 삭제되지 않은 항목들 중에서 특정 인덱스를 골라 삭제하는 함수가 필요하다.

코드

#include <iostream>
#include <algorithm>
#include <vector>
#include <cmath>
using namespace std;
using ll = long long;
int n, m, k;
vector<int> arr, seg;
int init(int n, int s, int e) {
    if (s == e) return seg[n] = 1;
    int mid = (s + e) / 2;
    int s1 = init(n * 2, s, mid);
    int s2 = init(n * 2 + 1, mid + 1, e);
    return seg[n] = s1 + s2;
}
int del(int n, int s, int e, int ps, int p) {
    // 현재 구간에 제거 대상 인덱스가 없는 경우 탐색 종료. 
    // 1. ps + seg[n] < p : 앞선 구간과 현재 구간의 모든 수를 합쳐도 제거해야할 인덱스 p에 미치지 못하는 경우
    // 2. ps >= p : 현재 구간의 수를 포함하지 않아도 이미 제거해야할 인덱스 p 이상인 경우
    if (ps + seg[n] < p || ps >= p) return 0;
    seg[n]--; // 현재 구간이 제거된 수를 포함하므로 개수 - 1
    if (s == e) return s; // 리프 노드인 경우 제거된 수 반환
    int mid = (s + e) / 2;
    int s2 = del(n * 2 + 1, mid + 1, e, ps + seg[n * 2], p); // 오른자식 탐색부터 해야함. seg[n * 2] 가 왼자식 탐색에서 변화될 수 있음.
    int s1 = del(n * 2, s, mid, ps, p); // 왼자식 탐색
    return s1 + s2;
}
int main()
{
    ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);
    cin >> n >> k; seg = vector<int>((int)pow(2, (int)ceil(log2(n)) + 1));
    init(1, 1, n);    
    cout << "<";
    bool first = true;
    int p = 1;
    for (int i = 1; i <= n; i++) {
        if (!first) {
            cout << ", ";
        }
        else {
            first = false;
        }
        p = (p + k - 2) % seg[1] + 1; // (현재 위치 p + k - 1 (인덱스 보정용 1) - 1 (현재 위치 p는 이미 삭제되었으므로 모든 배열은 한칸 당겨진 상태)) % 수열 길이 + 1 (인덱스 보정)
        cout << del(1, 1, n, 0, p);
    }
    cout << ">";
    return 0;
}