[JAVA] 19942 - 다이어트 💛4: 입력의 크기가 30 이하인 조합 문제는 비트마스킹으로 풀자!

들어가며

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

문제는 경우의 수를 결국 다 돌아봐야 하는 완전탐색 문제이다. 풀고 나서 다른 분들의 풀이를 보니까 비트마스킹 기법으로 푸신분들이 많아서 이번기회에 간단하게 정리하고 다음에도 써보고자 한다.

본문으로

일단은 재귀적으로 경우의 수를 탐색하면서 원상복귀해주는 느낌으로 아래처럼 구현했다. 재귀 함수의 인자가 매우 많지만,,, 풀려고 최선을 다했다.

public class Main {

    static int N;
    static int sProtein, sFat, sCarbo, sVitamin;
    static Food[] foods;
    static List<Integer> retList;
    static int ret;

    static class Food {
        int protein;
        int fat;
        int carbo;
        int vitamin;
        int price;

        Food(int protein, int fat, int carbo, int vitamin, int price) {
            this.protein = protein;
            this.fat = fat;
            this.carbo = carbo;
            this.vitamin = vitamin;
            this.price = price;
        }
    }

    static void go(int protein, int fat, int carbo, int vitamin, List<Integer> foodList, int price, int here) {
        if (protein >= sProtein && fat >= sFat && carbo >= sCarbo && vitamin >= sVitamin) {
            if (price < ret) {
                ret = price;
                retList = new ArrayList<>(foodList);
                return;
            }
        }

        if (price > ret) {
            return;
        }

        for (int i = here; i < N; i++) {
            int dProtein = protein + foods[i].protein;
            int dFat = fat + foods[i].fat;
            int dCarbo = carbo + foods[i].carbo;
            int dVitamin = vitamin + foods[i].vitamin;
            foodList.add(i);
            int dPrice = price + foods[i].price;
            go(dProtein, dFat, dCarbo, dVitamin, foodList, dPrice, i + 1);
            foodList.remove(foodList.size() - 1);
        }

    }

    public static void main(String[] args) throws IOException {
        BufferedReader bf = new BufferedReader(new InputStreamReader(System.in));
        N = Integer.parseInt(bf.readLine());
        StringTokenizer st = new StringTokenizer(bf.readLine());
        sProtein = Integer.parseInt(st.nextToken());
        sFat = Integer.parseInt(st.nextToken());
        sCarbo = Integer.parseInt(st.nextToken());
        sVitamin = Integer.parseInt(st.nextToken());

        foods = new Food[N];
        for (int i = 0; i < N; i++) {
            StringTokenizer line = new StringTokenizer(bf.readLine());
            int protein = Integer.parseInt(line.nextToken());
            int fat = Integer.parseInt(line.nextToken());
            int carbo = Integer.parseInt(line.nextToken());
            int vitamin = Integer.parseInt(line.nextToken());
            int price = Integer.parseInt(line.nextToken());
            foods[i] = new Food(protein, fat, carbo, vitamin, price);
        }

        ret = Integer.MAX_VALUE;
        List<Integer> foodList = new ArrayList<>();
        retList = new ArrayList<>();
        go(0, 0, 0, 0, foodList, 0, 0);

        if (ret == Integer.MAX_VALUE) {
            System.out.println(-1);
        } else {
            System.out.println(ret);
            for (int i = 0; i < retList.size(); i++) {
                System.out.print(retList.get(i) + 1 + " ");
            }
        }

    }
}

정답은 잘 나온다! 하지만 다른분들의 풀이를 참고하다 보면 <<나 &등의 비트연산자를 활용하여 풀이를 진행하시는 분들이 꽤 있었다. 사실 비트 연산자를 코딩테스트에 활용한 적은 한번도 없는데, 이번 기회에 조금 정리하며 다음번에는 비트마스킹으로 조합을 뽑아볼까 한다.

---

비트마스킹 뭐더라?

학부 1학년 때 OR, XOR 연산 비트 전환 등등 여러가지 배웠던 걸로 기억하는데,,, 물론 군대 다녀오니 까먹었다!! ㅎㅎ. 비트 연산자가 헷갈리는 사람들은 구글에 비트 연산자에 대해서 검색해보길 바란다. 비트를 다루기 위해 정의된 연산자일 뿐이다.

하지만 이를 문제에 적용하는 것은 어려웠다. 아래는 비트마스킹을 활용해서 조합을 뽑아내는 정답 코드이다.

public class Main {

    static int N;
    static int sProtein, sFat, sCarbo, sVitamin;
    static Food[] foods;

    static class Food {
        int protein, fat, carbo, vitamin, price;

        Food(int protein, int fat, int carbo, int vitamin, int price) {
            this.protein = protein;
            this.fat = fat;
            this.carbo = carbo;
            this.vitamin = vitamin;
            this.price = price;
        }
    }

    static boolean makeSort(List<Integer> curList, List<Integer> bestList) {
        int len = Math.min(curList.size(), bestList.size());
        for (int i = 0; i < len; i++) {
            if (!curList.get(i).equals(bestList.get(i))) {
                return curList.get(i) < bestList.get(i);
            }
        }

        return curList.size() < bestList.size();
    }

    public static void main(String[] args) throws Exception {
        BufferedReader bf = new BufferedReader(new InputStreamReader(System.in));

        N = Integer.parseInt(bf.readLine().trim());
        StringTokenizer st = new StringTokenizer(bf.readLine());
        sProtein = Integer.parseInt(st.nextToken());
        sFat = Integer.parseInt(st.nextToken());
        sCarbo = Integer.parseInt(st.nextToken());
        sVitamin = Integer.parseInt(st.nextToken());

        foods = new Food[N];
        for (int i = 0; i < N; i++) {
            st = new StringTokenizer(bf.readLine());
            int p = Integer.parseInt(st.nextToken());
            int f = Integer.parseInt(st.nextToken());
            int c = Integer.parseInt(st.nextToken());
            int v = Integer.parseInt(st.nextToken());
            int price = Integer.parseInt(st.nextToken());
            foods[i] = new Food(p, f, c, v, price);
        }

        int bestPrice = Integer.MAX_VALUE;
        List<Integer> bestList = new ArrayList<>();

        for (int i = 1; i < (1 << N); i++) {
            int p = 0, f = 0, c = 0, v = 0, price = 0;
            for (int j = 0; j < N; j++) {
                if ((i & (1 << j)) != 0) {
                    Food food = foods[j];
                    p += food.protein;
                    f += food.fat;
                    c += food.carbo;
                    v += food.vitamin;
                    price += food.price;
                }
            }

            if (p >= sProtein && f >= sFat && c >= sCarbo && v >= sVitamin) {
                List<Integer> curList = new ArrayList<>();
                for (int k = 0; k < N; k++) {
                    if ((i & (1 << k)) != 0) {
                        curList.add(k);
                    }
                }
                if (price < bestPrice) {
                    bestPrice = price;
                    bestList = curList;
                } else if (price == bestPrice) {
                    if (makeSort(curList, bestList)) {
                        bestList = curList;
                    }
                }
            }
        }

        if (bestPrice == Integer.MAX_VALUE) {
            System.out.println(-1);
            return;
        }

        StringBuilder sb = new StringBuilder();
        sb.append(bestPrice).append('\n');
        for (int x : bestList) {
            sb.append(x + 1).append(' ');
        }
        System.out.print(sb);
    }
}

---

아래의 로직이 비트마스킹을 활용해서 조합을 뽑아내는 핵심로직인데,, 

      for (int i = 1; i < (1 << N); i++) {
            int p = 0, f = 0, c = 0, v = 0, price = 0;
            for (int j = 0; j < N; j++) {
                if ((i & (1 << j)) != 0) {
                    Food food = foods[j];
                    p += food.protein;
                    f += food.fat;
                    c += food.carbo;
                    v += food.vitamin;
                    price += food.price;
                }
            }

 (1 << j) 는 j가 커짐에 따라 10, 100, 1000, 10000식으로 표현이 가능한데,  i를 증가시키면서(이진수로 생각하면, 01, 10, 11, 100...식으로 증가한다.) j번째 자릿수의 비트가 켜져있나 안켜져있나 확인하면 켜져있으면 그 자릿수의 음식을 추가한다. 그림으로 나타내면 아래처럼 나타낼 수 잇겠다. 순서대로 만들기 때문에 오름차순이 유지된다는 것도 장점이다.

시간복잡도도 크게 차이안난다!

 

마무리하며

이번에 비트마스킹이라는 기법을 남들 코드를 훔쳐보며 처음 배우게 되었는데, 입력값이 30이하일 때 조합을 뽑거나 상태를 관리하는 배열이 필요한 상황에서는 매우 유용한 기법인 것 같다! 그전에는 조합을 만드려면 몇개를 뽑는 로직을 따로 만들거나 방문 배열을 선언하는 등의 번거로움이 있었다. 비트마스킹을 활용하면 boolean 배열을 하나의 숫자로 표현 할 수 있는 것이 매우 큰 매력인 것 같다.

 

앞으로도 익숙해져서 애용해보자.