조합(Combination) 구현

이 글에서는 swift로 n개의 원소가 주어진 배열이 주어졌을 때 r개(0 < r <= n)의 결과를 구하는 코드를 구현해 보았습니다.

위 내용을 코드로 바꾸면 아래와 같습니다.

func uCombination<T>(elements: ArraySlice<T>, k: Int) -> [[T]] {

    if k == 0 {
        return [[]]
    }

    guard let first = elements.first else {
        return []
    }

    let head = [first]
    // n-1Cr-1
    let includeFirst = uCombination(elements: elements.dropFirst(), k: k - 1)
    // 첫 번째 원소는 반드시 포함한 조합의 결과
    var ret = includeFirst.map { head + $0}

    // 첫 번째 원소는 반드시 포함하지 않는 조합의 결과
    // n-1Cr
    var excludeFirst = uCombination(elements: elements.dropFirst(), k: k)

    return ret + excludeFirst
}

/// Array 타입에 대해서도 조합을 구현하기 위한 함수
func uCombination<T>(elements: [T], k: Int) -> [[T]] {
    return uCombination(elements: ArraySlice(elements), k: k)
}

이제 작성한 코드로 [1,2,3,4]로 주어진 배열에서 3개를 조합하는 코드를 작성하고 출력해 보겠습니다.

let arr = (1...4).map {$0}
let results = uCombination(elements: arr, k: 3)
results.forEach {print($0)}