itertools 모듈 combination 함수 사용하기

13594 단어 TILTIL

일단 되게 하라 말고 효율적으로!

알고리즘 공부중 내 수준에선 꽤 어려운 문제를 만났다. (물론 어디까지나 내 수준이다...)

문제의 조건 중 카드 3장을 뽑아야 하므로(당연히 중복으로 뽑으면 안 될 것이다.)
나는 3중 for문을 사용해서 구현해야겠다고 생각했다.


내가 짠 코드는 다음과 같다.

n, m = map(int, input().split())   # 카드의 개수 n  과 3장의 카드합으로 근접해야 하는 숫자 m

card_num = list(map(int, input().split()))   
# 카드의 개수 n장에 해당하는 카드숫자를 리스트로 만들어서 card_num 에 저장

result = 0        # 결과값을 저장할 result 라는 변수 선언

for i in range(n):           # 중복없이 3장의 카드 뽑기 위해 3중 for 문 사용
    for j in range(i+1, n):
        for k in range(j+1, n):
            if card_num[i] + card_num[j] + card_num[k] <= m:  
            # 뽑은 임의의 카드 3장의 합이 m보다 작거나 같다면 result 후보가 되므로
            # 더한 값을 result 에 저장
                result = max(result, card_num[i] + card_num[j] + card_num[k]) 
                # max 함수 사용해서 큰 값을 찾아감
                # 3중 for 문을 전부 돌고 나면 m의값에 가장 근접한 카드의 합이 result에 저장됨

print(result)

처음 문제를 풀고나서는 그럭저럭 코드가 엄청 길지도 않고 나름 잘짰다고 생각했다.
하지만 3중 for 문을 사용하면 시간복잡도가 상상이상으로 커질 수 있어 속도가 느려진다고 들었다.

더 효율적인 방법으로 문제를 해결 할 수 없을까 해서 구글링을했다.
(현재 내 수준으로는 3중 for문이 최선이었다....)


순열과 조합을 한방에 해결! itertools 모듈의 combination 함수!

A, B, C 라는 세 개의 문자가 있다고 가정하고

이 문자들중 두 개를 무작위로 뽑았을때의 순열과 조합은 어떻게 구하면 될까?


일단 경우의 수가 적으니 일일이 손으로 적어봐도 금방 나온다.

순열은 AB, AC, BA, BC, CA, CB 가 될 것이고,

조합은 AB, AC, BC 가 될 것이다.


but,!

실제로 (현업 또는 알고리즘 문제풀이) 일일이 손으로 적을 순 없을 것이다.
(누가 나를 개발자라고 불러줄 것인가.. 그냥 노가다꾼이지..)


이걸 가능하게 해주는게 파이썬 itertools 모듈의 permutations 과 combinations이다.

구구절절 설명보단 코드로!
(코드는 다음과 같다.)

import itertools

chars = ['A', 'B', 'C']

p = itertools.permutations(chars, 2)  # 순열
c = itertools.combinations(chars, 2)  # 조합

위 코드를 실행하게 되면,

p에는 itertools.permutations 객체가,

c에는 itertools.combinations 객체가 반환된다.


설명을 조금 덧붙이자면,

두 번째 인자로 받는 숫자(2)는 주어진 컨테이너 타입 변수에서 몇 개의 아이템을 조합할지 결정하는 인자이다.
그리고 permutations는 두 번째 인자를 받지 않으면 컨테이너의 전체 길이(위 예제에서는 3)가 default로 들어가고,
combinations는 두 번째 인자를 받지 않으면 동작하지 않는다(에러)

이를 list로 만들어서 출력해보면 다음과 같다.

import itertools

chars = ['A', 'B', 'C']

p = itertools.permutations(chars, 2)  # 순열
c = itertools.combinations(chars, 2)  # 조합

print(list(p))
print(list(c))


[('A', 'B'), ('A', 'C'), ('B', 'A'), ('B', 'C'), ('C', 'A'), ('C', 'B')]
[('A', 'B'), ('A', 'C'), ('B', 'C')]
# 순열과 조합이 정상적으로 출력되었다.

이 함수를 통해 위의 문제를 좀더 효율적인 알고리즘을 사용해서 구현하는 것이 가능하다.



itertools 모듈을 이용해서 다시 짜본 코드

from itertools import combinations

n, m = map(int, input().split())   # 카드의 개수 n  과 3장의 카드로 근접해야 하는 숫자 m

card_num = list(map(int, input().split()))        
# 카드의 개수 n장에 해당하는 카드숫자를 리스트로 만들어서 card_num 에 저장

result = 0      # 결과값을 저장 할 result 라는 변수 선언

for cards in combinations(card_num, 3):   
# combinations 함수 사용하여 card_num 리스트에서 중복없이 카드 3장을 뽑는다.
    temp_sum = sum(cards)         # 뽑은 3장의 카드의 합을 임시 합계 변수에 저장
    if temp_sum <= m:             # 임시 함계가 m보다 작거나 같다면 result 후보군이 되므로
        result = max(result, temp_sum)      # temp_sum을 result에 담아줌

print(result)

실행결과는 당연히 처음 문제를 푼 코드와 마찬가지로 정답을 도출 할 수 있다.
하지만 시간복잡도를 훨씬 낮춰 프로그램의 퍼포먼스를 높였다.


요약 : 항상 더 좋은 알고리즘은 없을까? 하는 의문 갖기, 의문만 갖지 말고 직접 검색해보기!

좋은 웹페이지 즐겨찾기