조합, itertools 없이 시간초과 문제 해결하기

정수 A와 B가 입력될 때 ,

A개의 원소에서 B개의 원소를 뽑는 조합에 대한 총 경우의 수를 출력하는 프로그램을 작성하세요.

 

이 문제에 대해서 다음과 같이 풀어볼 수 있다.

 

import sys
input=sys.stdin.readline

from itertools import combinations

if __name__=="__main__":
    A,B=map(int, input().split())
    print(len(list(combinations(range(A), B))))

 

그러나 이렇게 풀면 시간초과 문제를 겪게 된다.

결국 재귀를 이용하는 등 조합 함수를 직접 짜주어야 한다. 

 

방법 1 

import sys
sys.stdin=open('input.txt', 'r')
input=sys.stdin.readline

def comb(n, r):
    a = 1 
    b = 1
    if n == r or r==0:
        return 1
    for i in range(r):
        a *= (n-i) 
        b *= (r-i)
    return a//b

if __name__=="__main__":
    A,B=map(int, input().split())   # A개의 원소들에서 B개를 뽑기
    print(comb(A, B))

 

방법 2 

# DP로 푸는 법 
import sys
sys.stdin=open('input.txt', 'r')
input=sys.stdin.readline

if __name__=="__main__":
    A,B=map(int, input().split())   # A개의 원소들에서 B개를 뽑기
    dp = [[0]*(B+1) for _ in range(A+1)]

    for a in range(A+1):
        for j in range(min(a, B) + 1):
            if j == 0 or j == a:
                dp[a][j] = 1
            else:
                dp[a][j] = dp[a-1][j-1] + dp[a-1][j]
    
    print(dp[A][B])

 

방법 3 

# from math import comb 로 푸는 법 
import sys
input=sys.stdin.readline

from math import comb

# 입력 받기
A, B = map(int, input().split())

# A개의 원소에서 B개의 원소를 뽑는 조합의 경우의 수 계산
result = comb(A, B)

# 결과 출력
print(result)

 

꼭 기억하자....