python 3 : 프로그래머스 Level. 1 체육복

한참 삽질하면서 왜 안되는지 자가진단 하려고 짠 테스트 코드... 짜 놓고도 원인을 못 찾아서 문제였지만 그래도 comprehension 연습한 셈 치자.

import random

student_number = random.randint(2, 30)
lost_list = sorted(set(random.randint(1, student_number) for x in range(random.randint(1, student_number))))
reserve_list = sorted(set(random.randint(1, student_number) for x in range(random.randint(1, student_number))))

print("n      (%d)" % student_number)
print("lost   (%d) =" % len(lost_list), lost_list)
print("reserve(%d) =" % len(reserve_list), reserve_list)
print(solution(student_number, lost_list, reserve_list))

for 문의 범위 안에 있는 인자를 for 문 안에서 삭제해서 오류가 남. 한참 끙끙대다 if 문에서 reserve 조건을 처리해서 해결하나 싶었는데 학생 수만큼 도느라 lost에 [3, 4]가 있고 reserve에 [4]가 있을 때 3을 처리하면서 reserve의 4를 삭제해버려서 답이 틀리는 경우가 있었다. 그걸 보완해서 바로 제출한 게 아래 코드.

def solution(n, lost, reserve):
    for student in range(1, n + 1):
        if student in lost and student in reserve:
            lost.remove(student)
            reserve.remove(student)
        elif student in lost and student - 1 in reserve:
            lost.remove(student)
            reserve.remove(student - 1)
        elif student in lost and student + 1 in reserve:
            if student + 1 not in lost:
                lost.remove(student)
                reserve.remove(student + 1)
    answer = n - len(lost)

    return answer

그리고 좀더 깔끔하게 하고자 장 씨의 도움을 받은 게 아래 코드.

if를 두 번 쓴 이유는 조건이 위아래가 다르므로 구분하기 위함.

def solution(n, lost, reserve):
    for student in range(1, n + 1):
        if student not in lost:
            continue
        if student in reserve:
            lost.remove(student)
            reserve.remove(student)
        elif student - 1 in reserve:
            lost.remove(student)
            reserve.remove(student - 1)
        elif (student + 1 in reserve) and (student + 1 not in lost):
            lost.remove(student)
            reserve.remove(student + 1)
    answer = n - len(lost)

    return answer

for 문을 학생 수만큼 돌지 않아도 되므로 회전을 최소화하려고 lost로 돌게 만든 코드

마지막 elif 문처럼 여러 문이 중첩되어 있으면 구분하기 쉽게 괄호()를 넣어주는 게 좋다고 한다.

def solution(n, lost, reserve):
    recover_count = 0
    for student in lost:
        if student in reserve:
            recover_count += 1
            reserve.remove(student)
        elif student - 1 in reserve:
            recover_count += 1
            reserve.remove(student - 1)
        elif (student + 1 in reserve) and (student + 1 not in lost):
            recover_count += 1
            reserve.remove(student + 1)
            
    answer = n - (len(lost) + recover_count)
    return answer

집합은 교집합, 합집합, 차집합을 구할 수 있다. 교집합은 집합 사이에 &를 넣어 표기한다. 여기서 lost와 reserve를 교집합을 빼고 집합으로 for 문에 넘겼다. 이러면 .remove를 했을 때 list에서보다 처리가 빠르다.

def solution(n, lost, reserve):
    dupl = set(lost) & set(reserve)
    lost, reserve = set(lost) - dupl, set(reserve) - dupl
    answer = n - len(lost)
    
    for student in lost:
        if student - 1 in reserve:
            reserve.remove(student - 1)
            answer += 1
        elif student + 1 in reserve:
            reserve.remove(student + 1)
            answer += 1

    return answer

그 이유는 리스트에서 뭔가를 찾아서 지운다는 것은 아래처럼 for문으로 리스트 안에 있는 그 값을 찾아서 지우느라 시간이 걸리기 때문이다. 반면 사전형이나 집합은 hash이므로 바로 값을 찾아서 삭제할 수 있다. 그래서 처리 속도가 빨라지는 것이다.

_lost.remove(f)

def remove(f):
    for item in _lost:
        if item == f:
            del _lost[f]
            break