[Python] 프로그래머스 / level 1 / 그리디 / 체육복

https://school.programmers.co.kr/learn/courses/30/lessons/42862

프로그래머스

 

def solution(n, lost, reserve):
    student = len(lost)
    
    reserve_i, lost_j = 0, 0
    
    while reserve_i < len(reserve) and lost_j < len(lost):
        
        # 여벌 체육복을 가져온 학생이 체육복 도난 당한 경우
        if reserve[reserve_i] == lost[lost_j]:
            reserve_i += 1
            lost_j += 1
            student -= 1
        
        # -1 했는 학생이 있을 때 
        elif reserve[reserve_i] - 1 == lost[lost_j]:
            reserve_i += 1
            lost_j += 1
            student -= 1
        
        # +1 했는 학생이 있을 때
        elif reserve[reserve_i] + 1 == lost[lost_j]:
            reserve_i += 1
            lost_j += 1
            student -= 1
            
    return n - student

 

이 문제는 리트코드의 쿠키부여 문제와 비슷하다(내 생각).

그래서 다음과 같이 코드를 작성했지만 시간초과. 

 

리스트 2개 모두를 순환해 최악의 경우 시간복잡도 O(len(reserve)*len(lost))가 됨.

 

여분의 체육복을 가져온 학생이 도둑 맞은 경우를 먼저 뺀 리스트를 만들면 시간 복잡도 줄어든다.

 

 

def solution(n, lost, reserve):
    
    _reserve = [r for r in reserve if r not in lost]
    _lost = [l for l in lost if l not in reserve]
    
    _reserve.sort()
    _lost.sort()
    
    for r in _reserve:
        q  = r-1
        s = r+1
        if q in _lost:
            _lost.remove(q)
        elif s in _lost:
            _lost.remove(s)
            
    return n - len(_lost)