[BAEKJOON][python] 백준 온라인 저지 문제 풀이 - 28325번: 호숫가의 개미굴

문제 링크

문제
KOI 호숫가에 여러 개미가 모여 사는 개미굴이 있다. 개미굴은 둥근 호수의 둘레를 따라 $1$부터 $N$까지의 번호가 붙은 $N$개의 방이 차례대로 원형으로 배치되어 있으며, 모든 $i$ ($1 \le i \le N-1$)에 대해 $i$번째 방과 $i+1$번째 방이, 그리고 $N$번째 방과 $1$번째 방이 통로로 직접 연결되어 있는 형태였다.

하지만 여러 이유로 인해 각 방에서 몇 개의 쪽방이 갈라지기 시작해서, 현재는 모든 $i$ ($1 \le i \le N$)에 대해, 개미굴의 $i$번째 방에 $C_i$개의 쪽방이 통로로 직접 연결되어 있다. $i$번째 방과 연결된 쪽방은 $i$번째 방 이외에 다른 방과 통로로 연결되어 있지 않다.

예를 들어 $N = 7$이고, $C = [3,0,0,1,0,2,0]$인 개미굴은 아래 그림과 같은 형태이다.

개미굴의 각 방과 쪽방에는 최대 한 마리의 개미가 살 수 있다. 만약 통로로 직접 연결되어 있는 두 곳(방 또는 쪽방) 모두에 개미가 살고 있다면, 두 개미는 서로 불편해한다. 이러한 불편함을 방지하기 위해, 현재 개미굴의 각 통로가 직접 연결하는 두 곳 중 최대 한 곳에만 개미가 살고 있다.

개미들은 똑똑하기 때문에, 이 조건을 만족하는 하에 최대한 많은 수의 개미들이 현재 개미굴에 살고 있다고 한다. 현재 개미굴의 구조가 주어질 때, 개미굴에 살고 있는 개미의 수를 구하는 프로그램을 작성하라.

입력
첫 번째 줄에 정수 $N$이 주어진다.

두 번째 줄에 각 방과 연결된 쪽방의 개수를 의미하는 $N$개의 정수 $C_1, C_2, \cdots, C_N$이 공백으로 구분되어 주어진다.

출력
첫 번째 줄에 개미굴에 살고 있는 개미의 수를 출력한다.

소스코드

from sys import stdin
input = stdin.readline

n = int(input())
c = list(map(int, input().split()))

spare = set() # 쪽방이 있는 굴의 index
for i in range (n) :
    if c[i] != 0:
        spare.add(i)

if len(spare) == 0: # 쪽방이 아예 없는 경우
    print(n // 2)

else:
    memo = [-1] * n # 개미가 살지 않는 경우: -1, 쪽방이 있는 경우: 0, 개미가 사는 경우: 1
    count = 0 # 살 수 있는 개미의 수
    for i in range (n) :
        if i in spare:
            count += c[i] # 쪽방을 모두 채우는 것이 효율적이므로, 쪽방 수만큼 count 증가
            memo[i] = 0
    if n == 2:
        if len(spare) == 1:
            count += 1 # n이 2이고 쪽방이 한 굴에만 있는 경우는 쪽방 수 + 1(쪽방이 없는 굴) 출력
        print(count)
    else:
        index = 0 # 쪽방이 있는 굴의 첫 번째 index
        for i in range (n) :
            if i in spare:
                index = i
                break
        for i in range (index, n + index): # 쪽방이 있는 굴의 index부터 한 바퀴 돌기
            if memo[i % n] == -1: # 만약 현재 index에 개미가 살지 않는다면
                if memo[(i - 1) % n] != 1 and memo[(i + 1) % n] != 1: # 현재 index의 양옆에 쪽방이나 개미가 있지 않다면
                    memo[i % n] = 1 # 현재 index에 개미가 살 수 있음
                    count += 1
        print(count)

엄청 빙빙 돌아갔지만.. 간단하게 풀리는 문제였습니다.
핵심은 쪽방을 먼저 모두 채운 후, 쪽방이 있는 굴의 인덱스를 체크해 놓고
쪽방이 없는 굴들을 번갈아가며 채우는 것이었습니다.
근데 쪽방이 있는 굴의 인덱스를 기준으로 개미가 살 굴을 정해야 하기 때문에 쪽방이 있는 굴의 인덱스부터 탐색을 시작해야 합니다.