[Python] combinations

프로그래머에서 알고리즘을 풀다가 combinations에 대해서 알게 되었다.

내가 푼 "소수 만들기"문제는 리스트가 들어왔을 때, 리스트 안 세 수를 더한 값이 소수인 개수를 return 하는 문제였다.

# 풀이 1

def solution(nums):
    sum = []
    for i in range(len(nums)):
        for j in range(i+1, len(nums)):
            for k in range(j+1, len(nums)):
                a = nums[i] + nums[j] + nums[k]
                sum.append(a)
                
    count = 0
    for i in sum:
        total = 0
        for j in range(1, i+1):
            if i % j == 0:
                total += 1
        if total == 2:
            count += 1

    return count

처음 풀었을 때는 세 수를 nums에서 가져오기 위해서 세 번의 for 문을 사용했다.
길었던 첫 번째 부분을 combinations을 사용하면 짧게 줄일 수 있다.

# 풀이 2

from itertools import combinations

def solution(nums):
    count = 0
    for item in combinations(nums, 3):
        total = sum(item)
        for i in range(2, total):
            if total % i == 0:
                break
        else:
            count+=1
    return count


print(solution([1, 2, 3, 4]))

combinations는 리스트에서 몇 개의 수를 골라 그 조합을 튜플 형태로 보여준다.

예를 들어 solution([1, 2, 3, 4])가 들어왔다고 가정했을 때 item의 조합은 (1, 2, 3), (1, 2, 4), (1, 3, 4), (2, 3, 4)가 된다.

또한 처음에는 sum이라는 빈 리스트에 더한 값을 하나씩 append를 해주었지만, 이번에는 sum() 함수를 사용하여 for문 안에서 갱신해줬다.