[파이썬] 백준 5577: RBY팡! (골드2)

[파이썬] 백준 5577: RBY팡! (골드2)

 

5577번: RBY팡!

---

문제

---

풀이

0. 방향성 생각

누가봐도 스택문제. 입력 크기가 10000이라 O(N^2)으로도 비빌 수 있다.

태그에 브루트포스 딱지가 붙어있어서 입력이 터지는 모든 경우에 대해서 탐색을 진행했다.

 

1. 입력 받기, 터지는 후보 찾기

import sys
input = sys.stdin.readline

n = int(input())
arr = []
for i in range(n):
    arr.append(int(input()))

idx_list = []
for i in range(n-3):
    now = arr[i]
    table = {1:0,2:0,3:0}
    
    for j in range(i+1,i+4):
        table[arr[j]] += 1
    
    temp = [val for val in table.values() if val != 0]
    if temp == [3] and now not in temp :
        idx_list.append((i,arr[i+1]))
    elif (temp == [1,2] or temp == [2,1]) and table[now] == 2 :
        idx_list.append((i,now))

입력 받아주고 루프를 한 번 돌면서 터지는 경우인지 조사한다.

터지는 경우는 1222 같이 다음에 연속된 3개가 모두 같은경우, 1112 같이 탐색 값과 같은게 다음에 두 번 나오면 된다.

idx_list에 현재 인덱스 i와 탐색한 구간에서 선택해야할 색깔을 추가한다.

 

2. 함수 작성

def ball_stack(stack):
    while True :
        n = len(stack)
        pop_idx = []
        x = 0
        while x < n :
            nx = x+1
            while nx < n and stack[nx] == stack[x]: 
                nx += 1
            if nx-x >= 4:
                pop_idx.extend(range(x,nx))
            x = nx
        if not pop_idx:
            break
        
        pop_idx = set(pop_idx)
        new_stack = []
        for i in range(n):
            if i not in pop_idx:
                new_stack.append(stack[i])
        stack = new_stack
    return stack

제거해야할 인덱스를 pop_idx에 추가한다.

리스트를 순회하면서 현재 탐색한 인덱스 색깔이 다음 색깔과 같으면 탐색 범위를 늘려준다.

4개 이상인 경우 제거할 수 있으므로 pop_idx에 추가한다.

모두 제거하고 for문을 돌려서 pop_idx에 없으면 스택을 쌓는다.

3. 정답 출력

if idx_list :
    stack_len = []
    for idx,change in idx_list:
        stack = arr.copy()
        stack[idx:idx+4] = [change]*4
        stack = ball_stack(stack)       
        stack_len.append(len(stack))           
    print(min(stack_len))
    
else :
    print(n)

바꿔도 아무런 변화가 없는 경우 stack_len 함수가 비어서 min을 사용했을 시 오류가 발생한다.

이를 나눠주기 위해서 if else문 사용.

앞에서 구한 idx_list를 참고하여 stack의 부분을 바꿔내고 함수를 실행시킨다. 남은 개수를 stack_len에 저장한다.

 

 

---

전체코드

import sys
input = sys.stdin.readline

n = int(input())
arr = []
for i in range(n):
    arr.append(int(input()))

idx_list = []
for i in range(n-3):
    now = arr[i]
    table = {1:0,2:0,3:0}
    
    for j in range(i+1,i+4):
        table[arr[j]] += 1
    
    temp = [val for val in table.values() if val != 0]
    if temp == [3] and now not in temp :
        idx_list.append((i,arr[i+1]))
    elif (temp == [1,2] or temp == [2,1]) and table[now] == 2 :
        idx_list.append((i,now))

def ball_stack(stack):
    while True :
        n = len(stack)
        pop_idx = []
        x = 0
        while x < n :
            nx = x+1
            while nx < n and stack[nx] == stack[x]: 
                nx += 1
            if nx-x >= 4:
                pop_idx.extend(range(x,nx))
            x = nx
        if not pop_idx:
            break
        
        pop_idx = set(pop_idx)
        new_stack = []
        for i in range(n):
            if i not in pop_idx:
                new_stack.append(stack[i])
        stack = new_stack
    return stack

if idx_list :
    stack_len = []
    for idx,change in idx_list:
        stack = arr.copy()
        stack[idx:idx+4] = [change]*4
        stack = ball_stack(stack)       
        stack_len.append(len(stack))           
    print(min(stack_len))
    
else :
    print(n)

코멘트

첫 테스트 케이스가 입력의 거의 10000개정도 들어온듯. 통과 못해서 pypy3로 풀었음

arr[i] 탐색할 때 뒤에 3개 탐색하는거를 다 반복해서 비효율적인 부분도 있고 폭발 케이스 나눠서 구하는 부분도 비효율적이어서 이 부분 수정하면 python3로도 시간 통과 할 수 있을듯