[백준 5002] 도어맨 (Python)

www.acmicpc.net/problem/5002

5002번: 도어맨

 

문제풀이

N, arr : 입력값

m_cnt, w_cnt : 남자, 여자 수

 

현재 차례의 사람을 입장시킬 수 있으면 바로 입장, 그렇지 않다면 그 다음 사람을 먼저 입장시켜 보는 방식의 그리디 알고리즘으로 쉽게 해결할 수 있다.

 

1. 남자 여자 차이가 N보다 작을 경우에는 현재 차례의 성별을 바로 입장한다.

 

2. 남자 여자 차이가 N일 경우에는 다음 두 가지를 확인한다.

- 현재 사람을 입장시켜 보았을 때 남여 차이가 N보다 작아질 경우 => 현재 차례 사람 입장

- 그렇지 않을 경우, 다음 사람을 먼저 입장 시켜보기 => 남여 차이가 N보다 작아진다면 다음 사람 먼저 입장, 그렇지 않으면 종료

 

3. 1-2 과정을 반복한다.

 

코드

import sys

if __name__ == '__main__':
    N = int(input())
    arr = list(map(str, sys.stdin.readline().rstrip()))
    m_cnt, w_cnt = 0, 0
    ans = 0
    for i in range(len(arr)):
        # 남자 여자 차이가 N보다 작을 경우에는 바로 증가
        if abs(m_cnt - w_cnt) < N:
            if arr[i] == 'M': m_cnt += 1
            else: w_cnt += 1
            ans += 1
        # 남자 여자 차이가 N일 경우
        else:
            if arr[i] == 'M':
                # 남자 차례에서 남자를 증가시켰을 때 차이가 N보다 작으면 남자 증가
                if abs(m_cnt + 1 - w_cnt) < N:
                    m_cnt += 1
                    ans += 1
                else:
                    # 남자가 많아서 여자를 증가시켜야 하는 경우
                    # 현재 순서 다음에 여자이면 순서 바꿔서 여자를 입장
                    if i + 1 < len(arr) and arr[i + 1] == 'W':
                        arr[i], arr[i + 1] = arr[i + 1], arr[i]
                        w_cnt += 1
                        ans += 1
                    # 순서 바꿔도 불가능 하면 종료
                    else: break
            else:
                # 여자 차례에서 여자를 증가시켰을 때 차이가 N보다 작으면 여자 증가
                if abs(w_cnt + 1 - m_cnt) < N:
                    w_cnt += 1
                    ans += 1
                else:
                    # 여자가 많아서 남자를 증가시켜야 하는 경우
                    # 현재 순서 다음에 남자이면 순서 바꿔서 남자를 입장
                    if i + 1 < len(arr) and arr[i + 1] == 'M':
                        arr[i], arr[i + 1] = arr[i + 1], arr[i]
                        m_cnt += 1
                        ans += 1
                    # 순서 바꿔도 불가능하면 종료
                    else: break
    print(ans)

 

결과

 

문제풀이나 코드에서 이상한 부분 지적은 언제나 환영합니다!