[알고리즘] 주사위고르기

재밌는 문제였다.
개인적으로 깔끔한 문제라고 생각되는 문제들은 다 완탐으로 쉽게 풀리는 대신 시간 또는 메모리 초과를 특정 자료구조나 알고리즘으로 극복하는 문제들이 좀 깔끔하고 좋아보인다.

이것도 일단 완탐으로 접근할 수는 있어서 좋은 문제라고 생각한다.

문제는 생략

일단은 완탐으로 돌려봤다.

import java.util.*;

class Solution {
    static boolean[] selected;
    static int n, winCount, maxWinCount;
    static int[] answer;
    static int[][] dices;
    public int[] solution(int[][] dice) {
        n = dice.length;
        selected = new boolean[n];
        dices = dice;
        winCount = 0;
        maxWinCount = 0;
        
        selectDice(0, 0);
        
        return answer;
    }
    
    private static void selectDice(int index, int count) {
        if(count >= n / 2) {
            winCount = 0;
            countWin(0, 0, 0);
            if(winCount > maxWinCount) {
                maxWinCount = winCount;
                int[] result = new int[n / 2];
                int r = 0;
                for(int i=0; i<n; i++){
                    if(selected[i]) {
                        result[r] = i + 1;
                        r++;
                    }
                }
                answer = result;
            }
            return;
        }
        
        if(index >= n) return;
        
        selected[index] = true;
        selectDice(index + 1, count + 1);
        selected[index] = false;
        selectDice(index + 1, count);
    }
    
    private static void countWin(int index, int selectedSum, int unSelectedSum) {
        if(index >= n) {
            if(selectedSum > unSelectedSum) winCount++;
            return;
        }
        if(selected[index]) {
            for(int value: dices[index]) {
                countWin(index + 1, selectedSum + value, unSelectedSum);
            }
        }else {
            for(int value: dices[index]) {
                countWin(index + 1, selectedSum, unSelectedSum + value);
            }
        }
    }
    
    
}

문제 조건에 따르면 완탐시에는 최대 10C5(주사위 5개 중복없이 선택) * 10^10이라는 경우의 수를 탐색해야 해서 시간초과가 난다.

그래서 고민하다가 요즘 애용하는 TreeMap을 쓰면 쉽게 풀린다는 걸 알아냈다.

10C5로 주사위를 선택한 후, 각각의 주사위 합을 TreeMap<Sum, Count> 형태로 저장해줬다.

그러면? 선택된 TreeMap에서 큰 합부터 꺼내어 해당 합보다 작은 Key를 갖는 TreeMap을 분리하여 모든 값을 더해주면 된다.

TreeMap이 이진트리 기반으로 되어있어 headMap()을 통해 분리하는 시간복잡도는 logN이다.

import java.util.*;

class Solution {
    static boolean[] selected;
    static int n, winCount, maxWinCount;
    static int[] answer;
    static int[][] dices;
    static TreeMap<Integer, Integer> selectedGroup, unSelectedGroup;
    
    public int[] solution(int[][] dice) {
        n = dice.length;
        selected = new boolean[n];
        dices = dice;
        maxWinCount = 0;
        
        selectDice(0, 0);
        
        return answer;
    }
    
    private static void selectDice(int index, int count) {
        if(count >= n / 2) {
            winCount = 0;
            selectedGroup = new TreeMap<Integer, Integer>();
            unSelectedGroup = new TreeMap<Integer, Integer>();
            findEverySum(true, 0, 0, 0);
            findEverySum(false, 0, 0, 0);
            countResult();
            if(winCount > maxWinCount) {
                maxWinCount = winCount;
                int[] result = new int[n / 2];
                int r = 0;
                for(int i=0; i<n; i++){
                    if(selected[i]) {
                        result[r] = i + 1;
                        r++;
                    }
                }
                answer = result;
            }
            return;
        }
        
        if(index >= n) return;
        
        selected[index] = true;
        selectDice(index + 1, count + 1);
        selected[index] = false;
        selectDice(index + 1, count);
    }
    
    private static void countResult() {
        while(!selectedGroup.isEmpty()) {
            int sum = selectedGroup.lastKey();
            
            // 꺼낸 합보다 작은 Sum을 가진 맵
            SortedMap<Integer, Integer> lower = unSelectedGroup.headMap(sum);
            int selectCount = selectedGroup.get(sum);
            int unSelectCount = lower.values().stream()
                .mapToInt(Integer::intValue)
                .sum();
            
            winCount += selectCount * unSelectCount;
            selectedGroup.pollLastEntry();
            if(unSelectCount <= 0) break;
        }
    }
    
    private static void findEverySum(boolean isSelectedGroup, int index, int sum, int count) {
        if(count >= n / 2) {
            if(isSelectedGroup) {
                selectedGroup.put(sum, selectedGroup.getOrDefault(sum, 0) + 1);
            }else {
                unSelectedGroup.put(sum, unSelectedGroup.getOrDefault(sum, 0) + 1);
            }
            return;
        }
        
        if(isSelectedGroup) {
            if(selected[index]) {
                for(int value: dices[index]) {
                    findEverySum(isSelectedGroup, index + 1, sum + value, count + 1);
                }
            }else {
                findEverySum(isSelectedGroup, index + 1, sum, count);
            }
        }else {
            if(!selected[index]) {
                for(int value: dices[index]) {
                    findEverySum(isSelectedGroup, index + 1, sum + value, count + 1);
                }
            }else {
                findEverySum(isSelectedGroup, index + 1, sum, count);
            }
        }
    }
    
    
}```