Algorithm

[Algo] 백준 19699 소-난다JAVA

조핑구 2024. 5. 8. 13:36

문제 바로가기 : https://www.acmicpc.net/problem/19699

메모리 : 14232KB

시간 : 124ms

언어 : Java 11


코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.Set;
import java.util.StringTokenizer;
import java.util.TreeSet;

public class Main {
    static int N, M, sum;
    static int[] cow;
    static Set<Integer> anslist;
    static boolean[] prime, check;

    public static void main(String[] args) throws NumberFormatException, IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        // 소수 판별
        prime = new boolean[10001];
        prime[0] = true;
        prime[1] = true;
        for (int i = 2; i < 10001; i++) {
            for (int j = i * 2; j < 10001; j += i) {
                prime[j] = true;
            }
        }

        StringTokenizer st = new StringTokenizer(br.readLine());
        N = Integer.parseInt(st.nextToken());
        M = Integer.parseInt(st.nextToken());
        cow = new int[N];
        st = new StringTokenizer(br.readLine());
        for (int i = 0; i < N; i++) {
            cow[i] = Integer.parseInt(st.nextToken());
        }
        anslist = new TreeSet<>();
        check = new boolean[N];
        sum = 0;
        //조합으로 찾기
        com(0, 0);
        StringBuilder sb = new StringBuilder();
        if (anslist.isEmpty()) {
            System.out.println(-1);
        } else {
            for (int i : anslist) {
                sb.append(i).append(" ");
            }
            System.out.print(sb);
        }
    }

    private static void com(int idx, int cnt) {
        if (cnt == M) {
            if (!prime[sum]) {
                anslist.add(sum);
            }
        }
        if (idx == N)
            return;
        sum += cow[idx];
        check[idx] = true;
        com(idx + 1, cnt + 1);
        sum -= cow[idx];
        check[idx] = false;
        com(idx + 1, cnt);
    }
}

풀이

문제에는 두 가지 조건이 있다.

  1. N마리의 소 중 M마리를 골라야 한다.
  2. 선택된 M마리의 몸무게의 합이 소수여야 한다.

1번 조건은 조합으로, 2번 조건은 에라토스테네스의 체 로 해결할 수 있다.

소들의 몸무게가 같은 경우가 있을 수 있으므로 Set을 사용해 중복을 제거해주고, 오름차순으로 출력해야하니 TreeSet을 사용했다.