[SWEA 4837].[파이썬 S/W 문제해결 기본] 2일차 - 부분집합의 합

SW Expert Academy에서 학습용으로 문제를 가져왔습니다. 문제가 될 시 수정, 삭제하겠습니다.

https://swexpertacademy.com/main/main.do

문제 : 1부터 12까지의 숫자를 원소로 가진 집합 A가 있다. 집합 A의 부분 집합 중 N개의 원소를 갖고 있고, 원소의 합이 K인 부분집합의 개수를 출력하는 프로그램을 작성하시오.
해당하는 부분집합이 없는 경우 0을 출력한다. 모든 부분 집합을 만들어 답을 찾아도 된다.
예를 들어 N = 3, K = 6 경우, 부분집합은 { 1, 2, 3 } 경우 1가지가 존재한다.

[입력]

첫 줄에 테스트 케이스 개수 T가 주어진다.  ( 1 ≤ T ≤ 50 )
테스트 케이스 별로 부분집합 원소의 수 N과 부분 집합의 합 K가 여백을 두고 주어진다. ( 1 ≤ N ≤ 12, 1 ≤ K ≤ 100 )

[출력]

각 줄마다 "#T" (T는 테스트 케이스 번호)를 출력한 뒤, 답을 출력한다.

input

3
3 6
5 15
5 10

output

#1 1
#2 1
#3 0

code

numbers=[_ for _ in range(1,13)] #1부터 12까지 숫자의 리스트를 만들어준다.

lst=[] #빈 리스트 생성

for i in range(1 << len(numbers)): #비트연산자로 1*2^len(numbers)
    sub_lst=[] #빈 부분집합 리스트
    for j in range(len(numbers)): 
        if i & (1<<j): 
            sub_lst.append(numbers[j])
    lst.append(sub_lst) #부분집합을 만들어 빈 리스트에 부분집합을 넣어준다.



Test_case=int(input())
for t in range(1,Test_case+1):
    N,K=map(int,input().split())
    count=0

    for s in lst:
        if len(s) == N and sum(s) ==K: #리스트의 길이는 N, 리스트의 합은 K
            count+=1

    print('#{} {}'.format(t, count))

문제의 흐름을 정리해봅시다.

 

1.문제에서 요구한대로 1부터 12까지의 숫자 리스트를 형성해줍니다.
2.비트연산중 shift연산자를 사용하였습니다.

ex)
1 << 0 이면 1 (이진수)
1 << 1 이면 10 (이진수)
1 << 2 이면 100 (이진수)
1 << 4 10000 이고 2^4인 16이다.
즉 numbers = [1,2,3,4,5,6,7,8,9,10,11,12] 이므로 총 1<< 12 == 1000000000000(이진수) == 2^12 ==5048 만큼 돈다는 것입니다.

3.빈 부분집합 리스트를 만드는데 if i & (1 << j) 이 구문이 생소할 수 있습니다. 간단히 설명하자면

&는 AND를 찾는 비트연산자로, 예를 들어 A & B에서 A와 B의 2진수 형태에서 공통된 값을 10진수로 변환해 리턴해주는 역할을 합니다.

예시)
print(5 & 4)
=>4

위의 예시는 4를 리턴해줍니다. 5의 2진수 형태는 101이고, 4의 2진수 형태는 100입니다. 그러므로 5와 4의 교집합은 100이기 때문에 이것을 10진수로 변환해 4를 출력합니다.

 

즉 i와 2의 j승이 지정된 범위안에서 교집합이면 부분집합 리스트에 넣어줍니다. 그 후 부분집합리스트를 만들어준 리스트(lst)에 넣어줍니다.

 

이 리스트(lst)를 가지고 만약 리스트의 길이가 N이고 리스트의 합이 K이면 count에 1을 더해줍니다.

 

 

간단한 code

import itertools
Test_case = int(input())

numbers = [i for i in range(1,13)]
for t in range(1,Test_case+1):
    N,K = map(int,input().split())
    combi = list(itertools.combinations(numbers,N))
    count = 0
    for i in combi:
        if sum(i)==K:
            count+=1
    print('#{} {}'.format(t, count))

itertools함수를 사용하여 더 간단하게 pass할 수 있습니다.

itertools의 combinations을 이용하여 1부터 12까지의 리스트에서 N개를 뽑아 합이 0이면 count에 1을 더해주므로서 부분집합의 수를 구할 수 있습니다.