파이썬(Python) 리스트 for 루프 중(iteration) 특정 아이템 제거하기

리스트(list)를 for나 while로 iteration 하는 도중에 특정 아이템을 제거하고자 할때는 아래와 같이 list comprehension 을 사용하면 됩니다.

mylist = [x for x in mylist if check_if_valid(x)]

 

아래 예제 코드로 확인해 볼께요. 이 예제는, 리스트를 순회하면서 음수들을 제거하는 코드입니다. 

# 아이템 x가 0 또는 양수면 True, 음수면 False 반환
def check_if_valid(x):
    if x >=0:
        return True
    return False

mylist = [-1, 2, 10, -100, 0, 5, 7, -3]
mylist = [x for x in mylist if check_if_valid(x)]
print(mylist)
# [2, 10, 0, 5, 7]

 

위의 방법은 직관적이긴 하지만 한가지 문제가 있습니다. 바로 메모리 이슈입니다. 위의 방법에선 특정 아이템을 필터링하기 위해서 mylist 전체 복사본이 생성되어야 하기 때문입니다. mylist 크기가 작다면 문제 될 것이 없지만 리스트가 메모리에 꽉 찰 정도로 크다면 프로그램이 실행 중에 에러를 내며 종료 될 것입니다. 

 

메모리 문제를 해결하기 위해서는 위의 list comprehension을 generator로 바꿔주어야만 합니다. 

gen = (x for x in mylist if check_if_valid(x))

 

정말로 generator가 맞는지 확인 해볼께요. 

import types

isinstance(gen, types.GeneratorType)
# True

 

이제 mylist에 generator가 생성하는 아이템을 넣어주면 됩니다. 

mylist[:] = gen
print(mylist)
# [2, 10, 0, 5, 7]