[알고리즘] 그리디_체육복

💡 강의 참고 전 코드 (성공)

def solution(n, lost, reserve):
    answer = 0
    lost.sort()
    reserve.sort()

    # 1. 여벌 가져온 학생 중 도난당한 학생 체크
    reserve_list = []
    for i in reserve: # for문 안에 remove 함수 있으면 시간복잡도 N^2....
        if i in lost:
            lost.remove(i)
            continue
        else:
            reserve_list.append(i)
    
    # 2. 체육복 유무 알 수 있는 having 
    having = [True for _ in range(n+1)]

    for i in lost:
        having[i] = False

    for i in reserve_list:
        if i == 1:
            if having[i+1] == False:
                having[i+1] = True
        elif i == n:
            if having[i-1] == False:
                having[i-1] = True
        else:
            if having[i-1] == False:
                having[i-1] = True
            elif having[i+1] == False:
                having[i+1] = True

    having[0] = False
    answer = having.count(True)

    return answer

통과는 했지만 코드가 매우 길고, for문 안에 remove 함수가 있어 시간복잡도가 O(N^2)이 나온다.

💡 강사님 풀이

우선 이 문제는 그리디 알고리즘을 이용해서 푸는 문제이다.
그리디 : 현재의 선택이 마지막 해답의 최적성을 해치지 않을 때 사용

def solution(n, lost, reserve):
	u = [1] * (n + 2) # 체육복 정보 (편의성을 위해 앞, 뒤 체육복을 가진 학생 추가)
    for i in reserve:
    	u[i] += 1 # 여벌 가져온 학생 체육복 하나씩 더 추가
    for i in lost:
    	u[i] -= 1
    for i in range(1, n+1):
    	if u[i-1] == 0 and u[i] == 2: # 앞의 사람 빌려야 하는 입장이고 내가 빌려줄 수 있으면
        	u[i-1:i+1] = [1, 1] # 체육복을 빌려준다.
        elif u[i] == 2 and u[i+1] == 0: # 아니면 뒷사람 체크
        	u[i:i+2] = [1, 1] # 뒷사람에게 체육복을 빌려준다.
	return len([x for x in u[1:-1] if x > 0])

초기 값을 모두 1로 둔 후에, 여벌체육복이 있으면 +1, 잃어버렸으면 -1을 해준다. 그리고 앞, 뒤 순서로 체크하면서 현재 위치의 학생이 여벌이 있고, 앞, 뒤의 학생이 체육복이 없는 경우 체육복을 빌려주는 방식으로 for문을 돌린다. (각 상황에서의 선택이 마지막 해답의 최적성을 해치지 않음)
이렇게 하면 시간복잡도는 O(N)이 된다.