백준 5615 아파트 임대

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

2xy+x+y 에 해당하는 형태의 아파트가 될 수 있는 면적인지 아닌지 파악하는 문제이다. 2xy+x+y = S인 S가 존재한다고 하자. 2S+1 = 4xy+2x+2y+1 = (2x+1)(2y+1)로 나타낼 수 있다. 만약 2S+1이 소수가 된다면 (2x+1)(2y+1)로 인수분해가 불가능하기 때문에 존재할 수 없는 면적이라는 결론을 내릴 수 있다.

면적은 최대 2^31-1 까지 가질 수 있기 때문에 일반적인 소수 판정으로는 시간내에 처리할 수 없다. 따라서 밀러 라빈 알고리즘을 이용해서 해결하도록 하자.

let alist:[UInt64] = [2, 7, 61]
func powmod(x: UInt64, y: UInt64, m:UInt64) -> UInt64 {
    var X = x%m
    var iter = y
    var result:UInt64 = 1
    while iter != 0 {
        if iter%2 == 1 {
            result = (result*X)%m
        }
        X = (X*X)%m
        iter /= 2
    }
    return result
}

func miller(n: UInt64, a: UInt64) -> Bool {
    var d = n-1
    while d%2 == 0 {
        if powmod(x: a, y: d, m: n) == n-1 {
            return true
        }
        d /= 2
    }
    let tmp = powmod(x: a, y: d, m: n)
    return tmp == n-1 || tmp == 1
}

func isPrime(n: UInt64) -> Bool {
    if n <= 1 {
        return false
    }
    for a in alist {
        if !miller(n: n, a: a) {
            return false
        }
    }
    return true
}

let t = Int(readLine()!)!
var result = 0
for _ in 0..<t {
    let k = UInt64(readLine()!)!
    let C = 2*k+1
    if C == 2 || C == 7 || C == 61 {
        result += 1
    }
    else if isPrime(n: 2*k+1) {
        result += 1
    }
}
print(result)