본문 바로가기

IT 공부/코딩테스트

[프로그래머스 level1] 삼총사

해설

리스트 요소의 3개를 더했을 때 0이 되면, 삼총사라고 하며

삼총사가 될 수 있는 경우의 수를 구하는 문제이다.

 

그러면 코드를 한글로 쓰자면,

1. number의 요소중 조합으로 3개를 뽑아서 

2. 3개를 더했을 때 0이 나오면

3. count값을 증가시키고

4. return 으로 count값을 출력하면 된다.

 

1번에서 순수하게 코드를 짜는 방법을 그리다가.. 분명히 파이썬에 조합을 구하는 모듈이 있을거라고 확신하고 구글에 검색했더니 역시나 나왔다.

import itertools import combinations

itertools이라는 라이브러리를 import해서 combinations을 사용한다.

list_a = list(combinations(number, 3))

3개의 요소를 뽑는데, 순서 상관 없이 뽑는다 -> 조합, 그리고 이를 list로 만들어 list_a에 담는다.

for num in list_a:

    if (num[0] + num[1] + num[2]) == 0:

        count+=1

요소는 튜플로 3개씩 뽑히기 때문에, for문을 돌려 튜플 하나하나에 접근한다.

튜플에 접근하면, 튜플 0,1,2인덱스를 사용하여 값을 더하고 그 값이 0이면 count를 1씩 증가시킨다.

코드

from itertools import combinations
def solution(number):
    count = 0
    list_a = list(combinations(number, 3))
    for num in list_a:
        if (num[0] + num[1] + num[2]) == 0: #다른 사람 풀이 보니까 sum(num) == 0:
            count +=1
    return count

다른 사람 코드

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

3개의 요소만 가지는 리스트이니까

[i,j,k]의 값만 설정하면 된다.

예를 들어

리스트의 요소가 총 5개라고 한다면

인덱스의 경우는 총 10개 5*4*3/3*2

012

013

014

023

024

034

123

124

134

234

이를 규칙화 하면

i는 최대 올 수 있는 경우 리스트의 길이-2(뒤에 2개의 요소(j,k)가 와야하니까)

j는 최대 올 수 있는 경우 전꺼(i)에서 +1 ~ 리스트의 길이 -1(뒤에 1개의 요소(k)가 와야하니까)

k는 최대 올 수 있는 경우 전꺼(k)에서 +1 ~ 리스트의 길이

그래서 정해주고 if문으로 sum으로 i,k,j 인덱스에 해당하는 수들을 더해주면 된다.