[파이썬 알고리즘 Tip] 반복문 안에서는 리스트에 절대 원소를 추가/삭제하지 않는다

알고리즘 문제가 아주 쉽든, 어려든 간에 for 문으로 리스트의 원소를 순회하는 코드는 항상 쓸 수밖에 없다.

 

N = 100
arr = [i for _ in range(N)]
to_be_removed = [1, 5, 7, 10]

for item in arr:
	if item in to_be_removed:
    	# 조건에 따라 어떤 명령을 실행

 

 

그런데 여기서 치명적인 실수를 종종 할 때가 있다.

바로 순회하는 iterable 객체(리스트)의 원소를 직접 건드리는 것이다.

 

이런 간단한 상황을 생각해보자.

 

arr = [1, 2, 3, 6, 7, 9]

for item in arr:
	if item % 3 == 0:
    	arr.remove(3)

 

주어진 리스트(arr)에서 어떤 조건(3으로 나눈 나머지가 0)을 만족하면, 그 배열에서 해당 원소를 제거한다.

 

머릿속에서 구현한 잘못된 시나리오는 다음과 같다.

 

 

 

3을 리스트에서 제거하고 나면, 당연히 그 다음 원소인 6을 체크할 거라 기대한다.

하지만 실제 코드에서는 6을 건너뛰고 바로 7을 체크한다.

 

 

 

리스트를 순회하는 도중에 원소를 추가하거나 제거하면,

기존 인덱스에 변화가 생겨서 특정 원소의 순회를 건너뛰거나 중복해서 순회하는 경우가 발생할 수 있다.

 

 

 

 

 

실제 문제에 적용해보기

코드트리의 다음 문제(<초대장과 번호표>)를 풀면서, 처음에 이렇게 코드를 작성했다.

 

https://www.codetree.ai/missions/8/problems/invitation-and-number-tag

코드트리 | 코딩테스트 준비를 위한 알고리즘 정석

 

 

처음에 라이브러리(deque)를 불러오고 입력 데이터를 받는 과정은 생략했다.

 

# 1번은 초대장을 무조건 준다
ans = 1

# 초대받은 사람을 한 명씩 설정하면서 for loop을 실행
# 처음에는 1번 사람이 초대된 것이므로, 1번을 invited에 저장
invited_people = deque([1])

# 반복
# for ... :
while invited_people:
    invited = invited_people.pop()
    for group_set in groups:
        # discard method를 쓰면, 해당 set에 원소가 존재하지 않아도 에러 X
        group_set.discard(invited)
        
        # 만약 초대받은 사람을 제외한 사람이 그룹에 1명밖에 없다면, 그 사람을 다음 초대받은 사람으로 설정
        if len(group_set) == 1:
            invited_people.appendleft(group_set.pop())
            ans += 1
        
        # 만약 그룹에 이제 남은 사람이 없다면, 그룹을 삭제
        if not group_set:
            groups.remove(group_set)
        
print(ans)

 

 

위 코드에서도 groups라는 리스트를 for 문으로 순회하고 있는데

리스트의 원소이자 집합인 group_set이 비어있다면 더 이상 탐색할 필요가 없기 때문에 삭제해주었다.

코드를 조금 더 효율적으로 만들겠다고 원소를 삭제해준 건데 엉뚱한 결과가 발생했다.

즉 groups 리스트의 인덱스가 어긋나버리기 때문에 특정 group_set을 아예 확인하지 않게 되고, 결국 답이 달라진다.

 

이 코드는 어떻게 고칠 수 있을까?

 

    for group_set in groups:
        # group_set이 비어있다고 삭제하는 대신, for 문을 건너뛰는 continue문을 추가
        if not group_set:
            continue
        ...

 

group_set 원소를 삭제하는 대신 if - continue 문을 추가만 해주면 된다.

인덱스 순서는 그대로 유지하면서 비어있는 group_set을 직접 탐색할 연산도 건너뛰기 때문에 시간을 어느 정도 단축할 수 있다.

 

 

 

 

마무리하며

iterable 객체(리스트)를 순회하는 도중 원소를 추가, 삭제하면 인덱스 순서가 어긋나게 된다.

따라서 실수를 미연에 방지하기 위해선, iterable 객체의 원소를 순회하는 도중엔 아예 건드리지 않는 게 좋다.

 

하지만 보통 원소를 동적으로 추가, 삭제하고 싶은 욕망이 드는 건 연산 시간을 단축하고 싶기 때문이다.

이런 경우에는 if 문으로 충분히 조건을 걸어서 탐색 시간을 줄일 수 있다.