[BOJ / Swift & Python] 17626 - Four Squares

BOJ 17626번 Four Squares

17626번: Four Squares

* 한 게시글에 코드 블록 언어를 하나로 통일해야 해서 Python 코드는 하이라이트가 정상적이지 않습니다ㅠ

Intro

처음에는 Python으로 DP를 사용하여 풀이하려 하였으나 계속해서 시간 초과가 발생했습니다.

구글링을 통해 DP를 사용한 Python 풀이를 찾아보았으나 거의 동일한 코드를 돌려도 시간 초과가 발생했습니다.
(테스트 케이스가 최근 추가됐거나 발견하지 못한 부분에서 효율이 안 좋았던 것 같습니다.. 혹시 원인을 아시는 분 있다면 조언 부탁드립니다.)

결론적으로는 완전 탐색으로 풀어 해결했습니다.

Solution

Failure Solution

우선 DP를 사용하여 시간초과로 인해 실패한 풀이를 공유하겠습니다.

1. DP 리스트에 0, 1을 담아둡니다.
2. 2부터 n까지 for문을 돌리며 2부터 n까지 모든 최소 개수를 구해 DP 리스트에 저장합니다.
3. 마지막으로 DP[n]을 출력합니다.

import sys
n = int(sys.stdin.readline())
DP = [0, 1]
for i in range(2, n+1) :
    DP.append(min([DP[i-(j**2)]+1 for j in range(1, int(i**0.5)+1)]))
print(DP[n])

Success Solution

완전 탐색을 통해 차례로 답이 1일 경우, 2일 경우, 3일 경우, 4일 경우를 테스트합니다.

1. sqrt(num)이 정수라면 1을 출력
2. a^2 + b^2가 num이라면 2를 출력
3. a^2 + b^2 + c^2가 num이라면 3을 출력
4. 위 1 ~ 3에 해당하지 않을 경우 4를 출력

Code

Swift

import Foundation
extension Double {
    func isInteger() -> Bool {
        return Double(floor(self)) == Double(self)
    }
}
func fs(_ num: Double) -> Int {
    if sqrt(num).isInteger() { return 1 }
    let a: Int = Int(sqrt(num))
    for i in 0...a where sqrt(num-Double(i*i)).isInteger() { return 2 }
    for i in 0...a {
        for j in 0...Int(sqrt(num-Double(i*i))) where sqrt(num-Double(i*i+j*j)).isInteger() {
            return 3
        }
    }
    return 4
}
print(fs(Double(readLine()!)!))

Python

import sys
def check(n) :
    if (n**0.5).is_integer() : return 1
    for i in range(1, int(n**0.5)+1) :
        if (int(n-i*i)**0.5).is_integer() : return 2
    for i in range(1, int(n**0.5)+1) :
        for j in range(1, int((n-i*i)**0.5)+1) :
            if ((n-i*i-j*j)**0.5).is_integer() : return 3
    return 4

n = int(sys.stdin.readline())
print(check(n))

 

* 개인 블로그에 작성한 내용을 복사해왔습니다.