[Python/파이썬] 백준 17615번 볼 모으기

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

 

코드 (100점)

import sys
input = sys.stdin.readline

n = int(input().strip())
balls = input().strip()

red_right = (balls.rstrip('R')).count('R')
red_left = (balls.lstrip('R')).count('R')
blue_right = (balls.rstrip('B')).count('B')
blue_left = (balls.lstrip('B')).count('B')

print(min([red_right, red_left, blue_left, blue_right]))

 

빨간색과 파란색이 각각 왼쪽으로 몰았을 때와 오른쪽으로 몰았을 때의 이동 횟수를 모두 구해보고, 그 중 최솟값을 구하는 방법이다.

 

파이썬의 내장 함수인 `strip()`과 `count()`를 적극적으로 활용하여 푼 것이 인상적이다.

처음 풀이한 코드(14점)

import sys
input = sys.stdin.readline

n = int(input().strip())
balls = input().strip()


def cnt_ball(balls):
    reds = []
    blues = []
    red_sum = 0
    blue_sum = 0
    for i in range(n):
        if balls[i] == 'R':
            if len(reds) == 0 or balls[i-1] == "B":
                reds.append(1)
            else:
                reds[-1] += 1
            red_sum += 1
        else:
            if len(blues) == 0 or balls[i-1] == "B":
                blues.append(1)
            else:
                blues[-1] += 1
            blue_sum += 1
    return [reds, blues, red_sum, blue_sum]


reds, blues, red_sum, blue_sum = cnt_ball(balls)

if red_sum == 0 or blue_sum == 0:
    print(0)
    exit()
answer = int(1e9)
# 빨간색을 왼쪽으로 몰기
if balls[0] == 'R':
    answer = min(answer, red_sum-reds[0])
else:
    answer = min(answer, red_sum)
# 빨간색을 오른쪽으로 몰기
if balls[-1] == 'R':
    answer = min(answer, red_sum-reds[-1])
else:
    answer = min(answer, red_sum)
# 파란색을 왼쪽으로 몰기
if balls[0] == 'B':
    answer = min(answer, blue_sum-blues[0])
else:
    answer = min(answer, blue_sum)
# 파란색을 오른쪽으로 몰기
if balls[-1] == 'B':
    answer = min(answer, blue_sum-blues[-1])
else:
    answer = min(answer, blue_sum)
print(answer)

 

처음에는 이렇게 빨간색과 파란색 덩어리들을 직접 구해서 풀이하려고 했다.

 

정답 코드로 통과하긴 했지만, 결국 로직도 비슷하고, 시간 복잡도도 O(N)으로 같은데 정답 코드와 처음에 푼 코드의 시간 차이가 많이 발생하는 이유가 궁금해서 찾아봤다.

 

출처: https://googleyness.tistory.com/126

파이썬의 내장 함수는 C언어로 구현되어 있어서 같은 동작을 하는 코드라도 직접 작성한 파이썬 코드보다 내장 함수가 더 빠르다고 한다. 그래서 이런 결과가 발생했던 것 같다.

---

 

그리고 추가적으로 발견한건데, 이건 위의 문제와는 관련없지만 결과가 흥미로워서 갖고 왔다.

출처: https://www.acmicpc.net/board/view/125895내

전역 변수를 사용하도록 코드를 작성한 것보다 함수로 작성했을 때 더 속도가 빠르다는 글이었다. 

 

내가 작성했던 코드에서도 적용이 되는 것인지 확인하기 위해 시간이 오래 걸리는 가장 큰 원인으로 예상되는 부분을 함수로 작성해봤다.

def cnt_ball(balls):
    reds = []
    blues = []
    red_sum = 0
    blue_sum = 0
    for i in range(n):
        if balls[i] == 'R':
            if len(reds) == 0 or balls[i-1] == "B":
                reds.append(1)
            else:
                reds[-1] += 1
            red_sum += 1
        else:
            if len(blues) == 0 or balls[i-1] == "B":
                blues.append(1)
            else:
                blues[-1] += 1
            blue_sum += 1
    return [reds, blues, red_sum, blue_sum]


reds, blues, red_sum, blue_sum = cnt_ball(balls)

 

이렇게 변경한 뒤에 제출하니 아래 사진과 같은 결과를 얻었다.

 

위에서부터 1,2,3이라고 하면

1: 3번 코드를 함수를 사용하도록 변경한 것. 바로 위의 코드.

2: strip() 사용한 100점 코드

3: 처음에 푼 코드

 

같은 코드인데도 함수로 묶어준 1번과 묶어주지 않은 3번이 확연히 차이나는 것을 확인할 수 있었다. 

 

결론

파이썬의 최적화를 위해서는 내장 함수를 적극 활용하고, 웬만하면 전역 변수를 사용하지 말자.

🔗 참고 링크

- https://googleyness.tistory.com/126

- https://www.acmicpc.net/board/view/125895