nums에 있는 숫자들 중 서로 다른 3개를 골라 더했을 때 소수가 되는 경우의 개수를 반환하는 함수 구현하기
주어진 숫자 중 3개의 수를 더했을 때 소수가 되는 경우의 개수를 구하려고 합니다. 숫자들이 들어있는 배열 nums가 매개변수로 주어질 때, nums에 있는 숫자들 중 서로 다른 3개를 골라 더했을 때 소수가 되는 경우의 개수를 return 하도록 solution 함수를 완성해주세요.
nums | result |
[1,2,3,4] | 1 |
[1,2,7,6,4] | 4 |
function solution(nums) {
var answer = 0;
let sum = 0;
for(let i = 0; i< nums.length-2;i++){
for(j=i+1; j<nums.length-1;j++){
for(k=j+1; k<nums.length;k++){
sum = nums[i]+nums[j]+nums[k];
let count = 0;
for(let l = 1; l <= sum; l++){
if(sum % l === 0) count++;
}
if(count == 2) answer++;
}
}
}
return answer;
}
테스트는 통과했지만, nums의 원소가 최대 1000까지 가능하기 때문인지, 일부 테스트의 경우 완료까지 오랜 시간이 걸렸고, 그렇다면 개선할 점이 있다고 판단했다.
소수를 판별하는 부분에 대한 함수를 구현했다.
function solution(nums) {
let answer = 0;
// 소수인지 판별
function isPrime(num){
if(num <= 1) return false;
if(num <= 3) return true; // 2, 3 소수
if(num % 2 === 0 || num % 3 === 0) return false;
for(let i = 5; i*i <= num; i+=6){
if (num % i === 0 || num % (i+2) === 0) return false;
}
return true;
}
for(let i = 0; i< nums.length-2;i++){
for(j=i+1; j<nums.length-1;j++){
for(k=j+1; k<nums.length;k++){
const sum = nums[i]+nums[j]+nums[k];
if(isPrime(sum)){
answer ++;
}
}
}
}
return answer;
}
개선한 함수를 Python으로 구현하면 아래와 같다.
from itertools import combinations
def solution(nums):
answer = 0
def is_prime(num):
if num <= 1:
return False
if num <= 3:
return True
if num % 2 == 0 or num % 3 == 0:
return False
i = 5
while i * i <= num:
if num % i == 0 or num % (i+2) == 0:
return False
i += 6
return True
for combo in combinations(nums, 3):
sum_value = sum(combo)
if is_prime(sum_value):
answer += 1
return answer
파이썬의 itertools.combinations
을 사용하여 배열 nums
에서 가능한 모든 3개 요소의 조합을 생성하면, 중첩된 반복문을 사용하지 않고도 조합을 간단히 구현할 수 있다.