[Python-CodingTest] 6-4. 합이 같은 부분집합 (DFS: 깊이우선탐색)

합이 같은 부분집합

문제 정리

입력

6
1 3 5 6 7 10

처리 과정

1. 원소들을 리스트 형태로 입력받기
2. 리스트의 합(=total) 계산
3. DFS의 파라미터로 인덱스(=i)와 누적합(=sum) 전달
4. i가 n과 같지 않을 때, DFS에 다음 인덱스(=i+1)와 누적합을 전달
5. 이때 누적합은 원소로 사용될 때는 해당 lst 값을 누적하고, 사용되지 않을 때는 그대로 전달
6. n-1까지 계산한 뒤, i가 n이 되면, 누적합(=sum)과 total-sum이 같은지 확인

출력

YES

내 답

import sys
sys.stdin = open("./input/in4.txt", "rt")

def DFS(i): # i는 인덱스 번호
    sum=0
    if i==n: # 맨 마지막 요소 다음
        for k in range(n):
            if ch[k]==1:
                sum+=numbers[k][1]
        if sum==total-sum:
            print('YES') # YES가 4번 출력..🤔
    else:
        ch[i]=1 # 부분집합으로 사용
        DFS(i+1)
        ch[i]=0 # 부분집합으로 사용X
        DFS(i+1)

if __name__=="__main__":
    n=int(input())
 
    # 튜플 형태로 받기
    numbers=[(idx,val) for idx,val in enumerate(list(map(int,input().split())))]

    ch=[0]*(n+1)
    total=0
    for i in range(n):
        total+=numbers[i][1]

    DFS(0) # 파라미터는 인덱스 번호

YES가 4번 출력됨..

풀이 - 1

import sys
sys.stdin = open("./input/in4.txt", "rt")

def DFS(i,sum): # i는 인덱스 번호, sum은 원소의 누적 합
    if i==n:
        if sum==total-sum:
            print('YES')
            sys.exit(0) # 프로그램 종료
    else:
        DFS(i+1,sum+lst[i]) # lst[i]를 원소로 사용 -> sum에 누적
        DFS(i+1,sum) # lst[i]를 원소로 사용X -> sum은 그대로

if __name__=="__main__":
    n=int(input())
    lst=list(map(int,input().split()))
    total=sum(lst) # lst의 합 구하기

    DFS(0,0)
    print('NO') # sys.exit(0)으로 종료되지 않았다면 NO

풀이 - 2: 시간 복잡도를 줄여보자

DFS() 안에서 sum==total-sum이 성립하려면, sum이 total의 절반을 넘기면 안됨!
부분집합의 원소를 sum에 누적하다가 total의 절반을 넘으면 더이상 진행하지 않고 NO

import sys
sys.stdin = open("./input/in4.txt", "rt")

def DFS(i,sum): 
    if sum>total//2: # 시간복잡도를 줄이기 위한 조건 추가
        return
    if i==n:
        if sum==total-sum:
            print('YES')
            sys.exit(0) 
    else:
        DFS(i+1,sum+lst[i]) 
        DFS(i+1,sum) 

if __name__=="__main__":
    n=int(input())
    lst=list(map(int,input().split()))
    total=sum(lst)

    DFS(0,0)
    print('NO') 

정리

• 프로그램을 종료할 때는 sys.exit(0)

---

💛 개인 공부 기록용 블로그입니다. 👻

맨 위로 이동하기