[Python] 9017. 크로스 컨트리

문제

크로스 컨트리 달리기는 주자들이 자연적인 야외의 지형에 만들어진 코스를 달리는 운동 경기이다. 경주 코스는 일반적으로 4에서 12 킬로미터이며, 숲이나 넓은 땅을 통과하는 풀과 흙으로 된 지면과 언덕과 평평한 지형을 포함한다. 이 경기는 주자들의 개인성적을 매기고, 이를 가지고 팀의 점수를 계산한다. 

한 팀은 여섯 명의 선수로 구성되며, 팀 점수는 상위 네 명의 주자의 점수를 합하여 계산한다. 점수는 자격을 갖춘 팀의 주자들에게만 주어지며, 결승점을 통과한 순서대로 점수를 받는다. 이 점수를 더하여 가장 낮은 점수를 얻는 팀이 우승을 하게 된다. 여섯 명의 주자가 참가하지 못한 팀은 점수 계산에서 제외된다. 동점의 경우에는 다섯 번째 주자가 가장 빨리 들어온 팀이 우승하게 된다.

예를 들어, 다음의 표를 살펴보자.

 

A	B	C	C	A	C	B	D	A	A	C	A	C	C	A
1	n/a	2	3	4	5	n/a	n/a	6	7	8	9	10	11	12

팀 B 와 D 는 선수의 수가 여섯이 아니므로, 점수를 받을 수 없다. 팀 A 의 점수는 18 (1+4+6+7)이고, 팀 C 의 점수는 18 (2+3+5+8)이다. 이 경우 두 팀의 점수가 같으므로 다섯 번째로 결승점을 통과한 선수를 고려한다, A 팀의 다섯 번째 선수의 점수가 C 팀의 다섯 번째 선수의 점수보다 적으므로 A 팀이 우승팀이 된다.

모든 선수들의 등수가 주어질 때, 우승팀을 구하는 프로그램을 작성하라. 각 팀의 참가 선수가 여섯보다 작으면 그 팀은 점수 계산에서 제외됨을 주의하라. 여섯 명 보다 많은 선수가 참가하는 팀은 없고, 적어도 한 팀은 참가 선수가 여섯이며, 모든 선수는 끝까지 완주를 한다고 가정한다.

---

 

입력 데이터는 표준입력을 사용한다. 입력은 T 개의 테스트 케이스로 주어진다. 입력 파일의 첫 번째 줄에 테스트 케이스의 수를 나타내는 정수 T 가 주어진다. 두 번째 줄부터는 두 줄에 하나의 테스트 케이스에 해당하는 데이터가 주어진다. 각 테스트 케이스의 첫 번째 줄에는 하나의 정수 N (6 ≤ N ≤ 1,000)이 주어진다. 두 번째 줄에는 팀 번호를 나타내는 N 개의 정수 t1, t2, …, tN 이 공백을 사이에 두고 주어진다. 각 팀은 1 과 M(1 ≤ M ≤ 200)사이의 정수로 표현된다.

풀이

list와 set으로 전체 순위와 유효한 팀 집합을 따로 관리하였다. 1~4번째 주자 순위 합, 5번째 주자 순위, 현재 몇 번째 주자인지를 기록하는 딕셔너리를 각각 만들어서 전체 순위 list를 한 번 순회하며 각 정보들을 채워넣고, 최종적으로 min()을 사용하여 승리 팀을 출력하였다.

첫 번째 코드

import sys

input = sys.stdin.readline

t = int(input())
for _ in range(t):
    n = int(input())
    rank = list(map(int, input().split()))
    team = {tm for tm in set(rank) if rank.count(tm) >= 6}

    fourth = {value: 0 for value in team}  # 1~4순위 합 딕셔너리
    fifth = {value: 0 for value in team}  # 5번째 인덱스 딕셔너리
    count = {value: 0 for value in team}  # 몇번째인지 기록

    ranking = 1
    for r in rank:
        if r in team:  # 유효한 팀이면
            count[r] += 1
            if count[r] <= 4:
                fourth[r] += ranking
            elif count[r] == 5:
                fifth[r] = ranking
            ranking += 1

    min_sc = min(fourth.values())
    winners = [k for k, v in fourth.items() if v == min_sc]

    if len(winners) == 1:
        print(winners[0])
    else:
        winner = min(winners, key=lambda k: fifth[k])
        print(winner)

나름 가독성이 괜찮게 구현된 것 같다. 코드 길이도 길이지만 특히 실행 시간을 많이 개선시킬 수 있을 것 같아 그 방향으로 코드를 수정해 보았다.

 

1. list를 value로 가지는 dictionary 선언

: 기존에는 1~4번째 주자 순위 합, 5번째 주자 순위, 몇 번째 주자인지 이 세가지에 대하여 각각 dictionary를 만들었지만, 이를 하나로 통합한 list를 value로 가지는 dictionary를 선언하였다. 

 

2. 매번 .count() 호출 대신 Counter() 사용

: for문에서 매번 .count()를 호출하는 대신, Counter()로 list의 각 원소의 개수를 value에 저장하는 dictionary를 선언해 효율성을 높였다.

3. min(), lambda 식을 이용해 튜플 기반 정렬

: 기존의 복잡한 우승팀 결정 코드를 min()과 튜플 기반 lambda식 정렬로 최적화하여 한 줄로 축약하였다.

 

최종 코드

import sys
from collections import Counter

input = sys.stdin.readline

t = int(input())
for _ in range(t):
    n = int(input())
    rank = list(map(int, input().split()))
    counts = Counter(rank)
    team = {t for t, c in counts.items() if c == 6}

    scores = {}  # {}로 선언하면 빈 딕셔너리 생성 [1~4 합, 5 값, 몇번째]
    ranking = 1
    for r in rank:
        if r in team:  # 유효한 팀이면
            if r not in scores:
                scores[r] = [0, 0, 0]
            scores[r][2] += 1
            if scores[r][2] <= 4:
                scores[r][0] += ranking
            elif scores[r][2] == 5:
                scores[r][1] = ranking
            ranking += 1
    winner = min(scores.items(), key=lambda x: (x[1][0], x[1][1]))

    print(winner[0])

Counter() 사용과 리스트 딕셔너리 선언으로 메모리 사용량이 약간 늘어났지만 코드 길이와 실행 시간은 대폭 줄어들었다.lambda 식 정렬과 리스트 기반 딕셔너리, Counter() 등 아직 잘 모르는 기능들을 많이 사용한 문제였다. 계속 반복하면서 익숙해지는 수 밖에 없는 것 같다...

 

배운 것

• Counter(list): from collections import Counter 한 후 사용. list의 원소별 나오는 개수를 {key:value}로 만든 dictionary를 반환한다.
• {}: 안에 들어가는 값에 따라 set, dictionary 둘 다 될 수 있지만 {}로 선언할 경우 빈 dictionary가 된다.
• min(scores.items(), key=lambda x: (x[1][0], x[1][1])): scores의 item들을 전부 가져온 뒤, 튜플() 내부의 값들 기준으로 다중 정렬한다. 이때 딕셔너리의 경우 key=x[0], value=x[1]로 넘어가기 때문에 value에 접근하기 위해선 앞에 x[1]를 붙여야 함.