백준 3955 캔디 분배

https://www.acmicpc.net/problem/3955

캔디를 총 KX + 1개를 구매하려고 하는데 사탕 한 봉지당 C개의 사탕이 들어있다고 할 때 문제의 조건을 만족하는 사탕 봉지의 갯수를 구하는 문제이다. 이 때 K와 C는 정해져있고 X는 자연수이다.

식을 다시 정리해보면 KX + 1 = CY 라고 쓸 수 있다. -KX + CY = 1로 식을 다시 정리할 수 있고 우리는 이를 만족하는 X와 Y를 찾고싶다. 확장 유클리드 알고리즘의 형태를 생각해보면 AX + BY = gcd(A, B)라는 것을 기억할 것이다.

즉 A = K, B = C, gcd(A, B) = 1인 부정방정식의 해를 찾는 문제로 바꿀 수 있다. K앞의 음수 부호는 최종적으로 구한 X값에 -를 취해주기만 하면 되기 때문에 별 문제없이 풀 수 있다.

따라서 확장 유클리드 알고리즘을 활용해 gcd값이 1일 때 Y해를 찾도록 하자. 유의할 것은 사탕 봉지의 최대 갯수가 정해져있기 때문에 해당 갯수를 넘어설 경우 실패처리를 해주자. 또한 오는 사람이 1명일 경우와 사탕 봉지의 사탕갯수가 1개일 경우 X = 0 혹은 Y = 0인 결과가 나오기 때문에 예외처리를 해줘야한다.

func calc(a: CLongLong, b: CLongLong) -> CLongLong {
    var r1 = a, r2 = b, s1:CLongLong = 1, s2:CLongLong = 0, t1:CLongLong = 0, t2:CLongLong = 1
    while r2 != 0 {
        let q = r1/r2
        let r = r1-q*r2, s = s1-q*s2, t = t1-q*t2
        r1 = r2
        s1 = s2
        t1 = t2
        r2 = r
        s2 = s
        t2 = t
    }
    t1 = (t1%a + a) % a
    if r1 != 1 || t1 > 1_000_000_000 {
        return -1
    }
    return t1
}

let t = Int(readLine()!)!
let MAX:CLongLong = 1_000_000_000
for _ in 0..<t {
    let kc = readLine()!.split(separator: " ").map{CLongLong(String($0))!}
    let k = kc[0], c = kc[1]
    let result = calc(a: k, b: c)
    if c == 1 {
        if k+1 <= MAX {
            print(k+1)
        }
        else {
            print("IMPOSSIBLE")
        }
    }
    else if k == 1 {
        print(1)
    }
    else if result != -1 {
        print(result)
    }
    else {
        print("IMPOSSIBLE")
    }
}