[Python] 32387번 충전하기

 

 

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

---

 

24/10/06

 

 

 

KCPC 당시, 파이썬으로 검수했던 문제다. 다양한 풀이가 있어서 나름 연습해보기 적절한 문제라고 생각한다.

---

 

문제 접근 방식:

 

 

문제의 요구 사항은 다음과 같다.

 

크기가 $N$인 Bool 배열 $A$가 주어지고, 두 종류의 쿼리가 주어진다. 쿼리의 개수는 최대 $10^6$개이며, 배열의 최대 크기는 $10^6$이다. 또한, 배열의 초기 상태는 모두 False로 초기화되어있는 상황이다.

 

1번 쿼리는 인덱스 $i$가 주어진다.

 

$i$보다 크거나 같은 인덱스 $k$중에서, $A[k]$ = False로 되어있는 가장 작은 인덱스를 True로 바꾸는 것이다. 만약 이것이 불가능하다면 $-1$을 출력한다.

 

2번 쿼리는 인덱스 $i$가 주어진다.

 

$A[i]$=False라면 $-1$을 출력한다. $A[i]$=True라면 가장 최근몇 번째 쿼리에서 True로 바뀌었는지를 출력하고, $A[i]$를 False로 바꾼다.

 

 

나이브하게 생각해볼 수 있는 것은, 각 쿼리마다, 정말 정직하게 주어진 그대로를 수행하는 것이다.

 

이는 시간 초과가 난다.

 

1번 쿼리를 $10^6$개 날리는데, $0$번 인덱스만 잔뜩 날린다고 생각해보자.

 

$i$보다 크거나 같은 인덱스 중에서 False로 되어있는 가장 작은 인덱스를 선형적으로 찾으려다가 시간초과가 나게 될 것이다.

 

시간 복잡도는 $\mathcal{O}(NQ)$가 될 것이다.

 

 

그러면 자연스럽게, 1번 쿼리에서 False로 되어있는 가장 작은 인덱스를 "선형적"이 아니라, "로그"로 줄여버리면 되지 않을까 하는 생각이 들 것이다.

 

왜냐하면 False가 지속되다가 True로 바뀌는 "가장 작은 지점"을 찾기만 하면 되기 때문이다.

 

그러면 자연스럽게 생각할 수 있는 알고리즘은 이분 탐색이다.

 

근데 문제점은 뭐냐면, False로 지속되다가 True로 바뀌는 순간이 여러 번 생긴다는 것이다.

 

 

예를 들어, 배열 $A$가 다음과 같이 주어져 있는데, 내가 1번 쿼리의 인덱스를 $2$로 날린다고 생각해보자.

 

$$[T, T, T, F, F, F, F, T, T, F, F, F]$$

 

이 경우 중간에 껴있는 True 두 개가 내 이분 탐색을 방해한다.

 

 

그러면 어떻게 해야할까?

 

다행히도 이런 상황을 지원하는 자료구조가 있다.

 

C++기준으로는 set이다.

 

엥 갑자기 set이라고 할 수 있는데, 내가 C++을 자세히 공부한 건 아니지만 듣기로는 C++의 set 자료구조가 레드블랙 트리로 구현되있어서, set에서도 이분탐색을 할 수 있다.

 

다른 주요 언어들(자바나 러스트 등)도 set이 레드블랙트리로 되어있어서, 그냥 내장함수를 적용하면 풀린다.

 

C++ 기준, set.lower_bound를 사용하면 된다. 이는 set컨테이너에서 i보다 크거나 같은 가장 작은 원소를 반환해준다.

 

그래서 $i$번째 포트가 가능함을 담아두는 set을 만들고 lower_bound를 적용하며 1번 쿼리를 해결하면 된다.

 

2번 쿼리는 사실 큰 문제가 되지 않는다. 그냥 또 다른 배열 만들고 그 배열에 언제 꼽았는지를 저장해주면 된다.

 

 

불행하게도 파이썬은 set의 기본 구현이 해시맵이다.

 

그래서 이 방법으로는 풀 수 없고, 직접 레드블랙트리를 구현하던가, 다른 방법으로 풀어야 된다.

 

한 가지 방법은 제곱근 분할법인데, 배열 $A$를 적당히 루트개로 쪼개서 관리하는 방법이다.

 

또 한 가지 방법은 세그먼트 트리인데, 이건 세그먼트 트리를 잘 하면 적용할 수 있을 것이라 생각한다.

 

다행히, 누군가 파이썬에도 C++처럼 레드블랙트리로 구현한 Set을 만들어놨다.

 

SortedSet이 그것인데, 불행하게도 내장 모듈이 아니다.

 

인터넷을 뒤지면 깃헙에 누군가 구현해놓은 모듈이 있고, 그 코드를 복붙해서 SortedSet으로도 해결할 수 있다.

 

나는 SortedSet으로 해결해서 코드 길이가 9만 바이트정도 된다. 그래서 파이썬 코드는 생략하도록 하겠다.

 

궁금하면 SortedSet을 검색해보도록 하자.

---

아래는 내가 위의 접근 방식과 같이 작성한 파이썬 코드이다. 더보기를 누르면 확인할 수 있다.

생략