문제 출처 : 프로그래머스

문제는 여기에

문제

• 학생들은 각자 정수 번호를 가지고 있다.
• 학생 3명의 정수 번호를 더했을 때 0이 되면 3명의 학생은 삼총사라고 한다.
• 학생들의 번호를 나타내는 정수 번호 number가 매개변수로 주어질 때, 학생들 중 삼총사를 만들 수 있는 방법의 수를 return 하도록 solution 함수를 완성하세요

제한사항

• 3 <= number 의 길이 <= 13
• -1,000 <= number 의 각 원소 <= 1,000
• 서로 다른 학생의 정수 번호가 같을 수 있습니다.

입출력 예

접근

 처음에는 깊게 생각하지 않고 for문을 세 번 돌리면 모든 경우에 대해 합이 3이 되는지 검사를 할 수 있겠다고 생각을 했다.

그런데 시간 초과가 발생할 것 같아 좀 더 고민해보았는데 다른 방식으로 해결하기에는 어려움이 있어 for 반복문을 세 번 돌리는 코드를 작성했는데 정답이 되었다. :>

def solution(number):
    cnt = 0
    for i in range(len(number)-2):
        for j in range(i+1, len(number)-1):
            for k in range(j+1, len(number)):
                if number[i] + number[j] + number[k] == 0:
                    cnt += 1
    return cnt

고찰

이렇게 무지성으로 for문 돌리는 것은 배울 것이 없어 고등학교 수학시간에 배운 조합(Combination)을 사용하여 문제를 풀어볼 것이다.

조합을 사용하기 위해 itertools 라는 파이썬에서 반복되는 데이터를 처리하는 기능을 포함하고 있는 라이브러리의 

Combination 클래스를 사용할 것이다.

def solution(number):
    from itertools import Combination
    cnt = 0
    for i in Combination(number, 3):    # 순서에 상관없이 가능한 모든 3가지 조합을 i에 할당
        if sum(i) == 0:
            cnt += 1
    return cnt

삼총사 문제로부터 코딩테스트에서 사용하면 유용한 itertools 라이브러리의 Combination 클래스 사용 법을 터득했다!