문제 링크: https://www.acmicpc.net/problem/32373

Thumbnail: Photo by Susan Holt Simpson (Unsplash)

시간초과와 싸우며 유익한 고민을 했던 문제였습니다 ㅎㅎ

풀이

어떻게 정렬을 할 것인가?

문제의 조건에서,  주어진 양의 정수 \(k\)에 대하여 수열에서 거리가 \(k\)인 쌍만 교환할 수 있다고 명시했습니다. 따라서, 일반적인 정렬 방법으로는 조건이 지켜지지 않습니다.

이런 경우에 사용할 수 있는 정렬 방법으로, 버킷 정렬(Bucket Sort) 가 있습니다.
간단히 설명하면, 적절한 개수의 버킷(배열)을 만들고 기존 배열의 원소들을 기준에 따라 각 버킷에 적절히 분배한 뒤 버킷을 다시 합치는 방법입니다.

이 문제에는 \(k\)개의 버킷을 만들고, 각 원소의 위치를 \(i + n \times k\) (단, \(n\)은 임의의 음이 아닌 정수) 나타냈을 때 \(i\)에 해당하는 버킷에 저장했습니다.

다시 말해, 현재 원소의 위치를 \(i\)라 하면 \(i \mod k\) 버킷에 원소를 저장합니다.

어떻게 버킷에 원소를 채울 것인가?

버킷에 원소를 채운 뒤, 이후에 다시 합치는 과정을 거쳐야 합니다.
하지만, 버킷에 원소를 정렬되지 않은 상태로 저장한다면 이후 합치는 과정의 시간 복잡도가 커지고, 결국 버킷 정렬을 택한 이유가 사라집니다.

단순히 각 버킷에 원소를 추가할 때마다 버킷을 순회하며 원소를 삽입할 위치를 찾아도 되지만, 이 경우도 시간 복잡도가 \(O(N^2)\)꼴이 되므로 적절하지 않습니다. 따라서, 이분탐색을 응용해 버킷에 원소를 추가하는 과정을 \(O(\log N)\)꼴로 개선해 보겠습니다.

계속해서 이분탐색으로 버킷에 원소를 추가한다면, 버킷은 정렬된 상태일 것입니다.
이제 이분탐색으로 원소를 찾는 과정을 생각해 봅시다.

이전에 장난감 경주 문제에서 공부한 '최적화 문제'의 응용으로, 결정 문제를 다음과 같이 정의하면 됩니다.

버킷에 새로 추가할 원소 \(e\)와 현재 버킷의 각 원소들의 위치 \(i\) (단, \(M\)은 현재 버키의 크기이며 \(0 \le i \le M\) 에 대해

• 버킷의 \(i\)번째 원소가 \(e\)보다 같거나 큰가?

이 결정 문제로 버킷의 각 원소들을 NNNN...NNYYY... 꼴로 만들고, 가장 먼저 나오는 Y의 앞에 원소를 추가해 주면 됩니다.

input = open(0).readline
N, K = map(int, input().split())
arr = list(map(int, input().split()))
buckets = [[] for _ in range(K)]

for i in range(N):
    elem = arr[i]
    bucket = buckets[i % K]

    l = 0
    r = len(bucket)
    # 결정 문제: elem이 bucket[m]보다 같거나 큰가?
    while l < r:
        m = (l + r) >> 1
        if bucket[m] < elem:    # N
            l = m + 1
        else:  # Y
            r = m
    bucket.insert(l, elem)

여담으로, 파이썬에는 내장 모듈 bisect에서 이분 탐색을 이용해 위와 같은 작업을 쉽게 수행할 수 있도록 내장 함수들을 제공합니다.

• bisect_left(arr, x): 순회 가능한(Iterable) 객체 arr에서 arr의 정렬 상태를 유지한 채로 x를 삽입할 수 있는 위치를 반환합니다. 이때, x와 같은 값이 arr안에 존재한다면 그 중 가장 왼쪽의 위치를 반환합니다.
  • lo 및 hi 매개변수로 사용할 배열의 범위를 지정할 수 있습니다. (선택)
  • key 매개변수로 원소의 크기를 비교할 때 사용할 값을 지정할 수 있습니다.
    두 원소 x와 y를 비교할 때, key(x)와 key(y)를 비교하게 됩니다. (선택)
• insort_left(arr, x): 순회 가능한(Iterable) 객체 arr에 x를 적절한 위치에 삽입합니다. 내부적으로는 bisect_left를 이용하며, 매개변수 또한 동일합니다.
• bisect_right(arr, x): 순회 가능한(Iterable) 객체 arr에서 arr의 정렬 상태를 유지한 채로 x를 삽입할 수 있는 위치를 반환합니다. 이때, x와 같은 값이 arr안에 존재한다면 그 중 가장 오른쪽의 위치를 반환합니다.
  • 매개변수들은 bisect_left와 동일합니다.

버킷 합치기

​이제, 앞서 정렬된 상태로 분할해 둔 버킷을 다시 합치면 됩니다. 원소의 위치에 따라 \(k\)개의 버킷으로 나누었으므로, 0번 버킷부터 \(k\)번 버킷까지 순차적으로 1개씩 원소를 꺼내 기존 배열로 합치면 됩니다.

이후, 합쳐진 배열이 여전히 정렬되어 있지 않다면 문제의 조건대로 정렬할 수 없는 배열이므로 "No"를, 정렬되었다면 "Yes"를 출력해 주면 됩니다.

arr[0] = buckets[0][0]
for i in range(1, N):  # 정렬되어있는지 판단.
    arr[i] = buckets[i % K][i // K]
    if arr[i] < arr[i - 1]:
        print("No")
        break
else:
    print("Yes")

다른 언어를 사용하시는 분들은 for~else가 익숙하지 않을 겁니다.
파이썬에서, 반복문 이후의 else문은 반복문이 break로 조기 종료되지 않았을 때에만 실행됩니다.

전체 코드

from bisect import insort_left
input = open(0).readline
N, K = map(int, input().split())
arr = list(map(int, input().split()))
buckets = [[] for _ in range(K)]

for i in range(N):
    elem = arr[i]
    bucket = buckets[i % K]

    # insort_left(bucket, elem) # 내장 구현 사용하기
    l = 0
    r = len(bucket)
    # 결정 문제: elem이 bucket[m]보다 같거나 큰가?
    while l < r:
        m = (l + r) >> 1
        if bucket[m] < elem:    # N
            l = m + 1
        else:  # Y
            r = m
    bucket.insert(l, elem)

    # print(f"bucket[{i%K}] = {bucket}")

arr[0] = buckets[0][0]
for i in range(1, N):  # 정렬되어있는지 판단.
    arr[i] = buckets[i % K][i // K]
    if arr[i] < arr[i - 1]:
        print("No")
        break
else:
    print("Yes")