[Algorithm] 프로그래머스 - 시소 짝궁 (Swift)

문제 설명

프로그래머스

시소 짝궁을 구하는 문제인데, 시소의 길이는 2, 3, 4가 있다. 만약 무게가 [100,180,360,100,270] 라고 한다면 [100, 100]은 시소 짝궁이 되고, [180, 360]도 4, 2길이에 앉게된다면 시소 짝궁이 된다.

 

문제 풀이

무게의 배열이 10만개가 주어지기 때문에 이중 for문을 하기에는 시간초과가 확실했다. 어떻게 하면 for문을 중복해서 사용하지 않을 수 있을까? 생각보다 정말 간단하게 접근이 가능했다. 우선 same 배열, multi 배열을 만들어서 무게의 최댓값이 1000이기 때문에 same 배열은 1001개 multi 배열은 4001개를 만들어서 index를 통해 각 무게를 계산하여 횟수를 넣었다. 그리고 핵심은 나중에 정답에 나누기 2를 하는 것이다.

import Foundation

func solution(_ weights:[Int]) -> Int64 {
    var answer = 0
    var same = Array(repeating: 0, count: 1001)
    var multi = Array(repeating: 0, count: 4001)
    
    for weight in weights {
        let m = weight
        let m2 = weight*2
        let m3 = weight*3
        let m4 = weight*4
        
        same[m] += 1
        multi[m2] += 1
        multi[m3] += 1
        multi[m4] += 1
    }
    
    for weight in weights {
        let m = weight
        let m2 = weight*2
        let m3 = weight*3
        let m4 = weight*4
        
        answer += same[m]-1
        answer += multi[m2]-same[m]
        answer += multi[m3]-same[m]
        answer += multi[m4]-same[m]
    }
    
    return Int64(answer/2)
}