[알고리즘, 코드트리] 색깔 트리 - java

문제

문제가 상당히 길다. 전문을 보고 싶은 사람은 코드트리에서 확인하자.

코드트리 | 코딩테스트 준비를 위한 알고리즘 정석

 

요약하자면 이렇다. 총 네 가지 동작을 수행해야 한다.

(1) 노드 추가
 - 노드를 트리에 추가합니다. 각 노드는 고유한 번호 mid​, 부모 노드 번호 pid​, 색깔 color 그리고 최대 깊이 maxdepth​ 를 가집니다.
 - 만약 기존 노드들의 maxdepth​ 값으로 인해 새로운 노드가 추가됨으로써 모순이 발생한다면, 현재 노드는 추가하지 않음에 유의합니다.

(2) 색깔 변경
 - 특정 노드 mid​ 를 루트로 하는 서브트리의 모든 노드의 색깔을 지정된 색 color 로 변경합니다.

(3) 색깔 조회
 - 특정 노드 mid​ 의 현재 색깔을 조회합니다.

(4) 점수 조회
 - 각 노드의 가치는 해당 노드를 루트로 하는 서브트리 내 서로 다른 색깔의 수로 정의됩니다. 모든 노드의 가치를 계산하여, 가치 제곱의 합을 출력합니다.

입력

첫 번째 줄에 명령의 수 Q 가 주어집니다.

두 번째 줄부터는 Q 개의 줄에 걸쳐 명령의 정보가 주어집니다. 각 명령에 따른 형태는 다음과 같습니다.

노드 추가

• 100 m_id p_id color max_depth 형태로 공백을 사이에 두고 주어집니다. 이는 고유한 번호 mid​, 부모 노드 번호 pid​, 색깔 color 그리고 최대 깊이 maxdepth​ 를 가지는 노드를 추가하는 연산이 주어졌음을 의미합니다.
• mid​ 는 새로 추가되는 노드 번호로 항상 새로운 값이 주어지며, pid​ 는 −1 이 아닌 이상 항상 이미 주어진 노드 번호가 주어짐을 가정해도 좋습니다.
• 해당 명령은 최대 20,000번 주어집니다.

색깔 변경

• 200 m_id color 형태로 공백을 사이에 두고 주어집니다. 이는 mid​ 노드를 루트로 하는 서브 트리 내 모든 노드의 색깔을 color 로 변경해야 함을 의미합니다. 이때 mid​ 는 이미 주어진 노드 번호만 주어짐을 가정해도 좋습니다.
• 해당 명령은 최대 50,000번 주어집니다.

색깔 조회

• 300 m_id 형태로 공백을 사이에 두고 주어집니다. 이 경우 현재 mid​ 노드는 어떤 색인지를 출력합니다. 이때 mid​ 는 이미 주어진 노드 번호만 주어짐을 가정해도 좋습니다.
• 해당 명령은 최대 20,000번 주어집니다.

점수 조회

• 400 형태로 주어집니다. 이 경우 모든 노드의 가치를 계산하여, 가치 제곱의 합을 출력합니다.
• 단, 이 연산은 노드가 1개 이상인 경우에만 호출되며 해당 명령은 최대 100번 주어집니다.

1≤Q≤100,000
1≤mid​, pid​≤100,000 (단, pid​ 는 −1 일 수 있음)
1≤color≤5
1≤maxdepth​≤100

출력

각 300과 400 명령에 대해 결과를 한 줄에 하나씩 순서대로 출력합니다.

문제 풀이

삼성 기출은 정말 "지문에 답이 있다". 기능 단위로 쪼개서 차분히 구현해보자.

자료 구조 정의 및 입출력

static class Node {
    int mId;
    int color;
    int maxDepth;
    Node parent = null;
    ArrayList<Node> children = new ArrayList<>();

    Node(int mId, int color, int maxDepth) {
        this.mId = mId;
        this.color = color;
        this.maxDepth = maxDepth;
    }

    void setParent(Node node) {
        this.parent = node;
    }

    void addChild(Node node) {
        this.children.add(node);
    }

    ArrayList<Node> getChildren() {
        return children;
    }

    void setColor(int color) {
        this.color = color;
    }

    int getColor() {
        return this.color;
    }
}

static ArrayList<Node> rootArr = new ArrayList<>();
static Map<Integer, Node> nodeMap = new HashMap<>();
static int result = 0;

int Q = Integer.parseInt(br.readLine());

while(Q-- > 0) {
    st = new StringTokenizer(br.readLine());

    int command = Integer.parseInt(st.nextToken());
    switch (command) {
        case 100:
            // addNode();
            break;
        case 200:
            // changeColor();
            break;
        case 300:
            // color = queryColor();
            // sb.append(color).append("\n");
            break;
        case 400:
            // score = queryScore();
            // sb.append(score).append("\n");
            break;
    }
}

System.out.println(sb);

전역적으로 추후에 있을 빠른 연산을 위해서 적절한 자료구조와 저장할 데이터의 종류를 정의해둔다. 기본적인 틀을 잡아 두었으니 이제 뭘 구현해야 할지 천천히 보이기 시작할 것이다.

노드 추가_ addNode()

static void addNode(Node node, int pId) {
    if(pId == -1) {
        rootArr.add(node);
        nodeMap.put(node.mId, node);
        return;
    }

    // maxDepth 검증 백트래킹
    Node parent = nodeMap.get(pId);
    boolean enableFlag = check(parent, 2);
    if(enableFlag) {

        // 부모 자식 관계 설정
        node.setParent(parent);
        parent.addChild(node);

        // 빠른 검색을 위해 map에 저장
        nodeMap.put(node.mId, node);
    } 
}

지문에 pId가 -1인 노드는 새로운 루트 노드라 정의했으므로 루트를 저장하는 rootArr 전역 배열과 mId로 특정 노드를 검색할 nodeMap 전역 Map 자료구조에 추가해준다.

루트 노드가 아니라면 pId를 통해 parent 노드를 찾고 트리를 거슬러 올라가면서 maxDepth를 검증한다.

직접 그렸다. 정성을 보았다면 공감을 누르자.

위 경우 maxDepth가 2인 살구색 노드에서 백트래킹으로 거슬러 올라온 최대 깊이가 maxDepth보다 크다면 해당 노드는 트리에 추가되어선 안된다. 이를 위한 검증함수 check()를 만들자.

maxDepth 검증함수 _ check()

static private boolean check(Node node, int depth) {
    // 루트일 때
    if(node.parent == null) {
        if(node.maxDepth < depth) {
            return false;
        } else {
            return true;
        }
    }

    boolean enableFlag = check(node.parent, depth + 1);

    if(enableFlag) {
        if(node.maxDepth < depth) {
            return false;
        } else {
            return true;
        }
    } else {
        return false;
    }
}

색깔 변경 _ changeColor()

static void changeColor(int mId, int color) {
    Node root = nodeMap.get(mId);
    root.setColor(color);

    // traversal
    traversal(root, color);
}

노드 추가에 비한다면 색깔 변경은 일도 아니다. 트리의 순회 방법만 알면 된다. 주어진 노드를 루트노드로 하는 서브 트리의 모든 색상을 바꾸면 되므로 간단하게 주어진 노드부터 순회하자.

트리 순회 _ traversal()

private static void traversal(Node node, int color) {

    for(Node child : node.getChildren()) {
        child.setColor(color);

        traversal(child, color);
    }
}

색깔 조회 _ queryColor()

static int queryColor(int mId) {
    return nodeMap.get(mId).getColor();
}

세번째 색깔 조회를 쉽게 구현하기 위해서 노드를 추가할 때마다 nodeMap 에 mId를 기준으로 모두 저장했다. 단순히 꺼내서 color를 가져오자.

점수 조회 _ queryScore()

static int queryScore() {
    result = 0;

    for(Node root : rootArr) {
        traversal2(root);
    }

    return result;
}

포레스트(여러 개의 트리)의 모든 트리를 방문하며 일일히 점수 계산을 해주어야 한다. 본인의 서브트리에 어떤 색상 조합이 있는지 기억해야 하므로 Set<Integer>을 반환하는 순회 함수 traversal2를 정의한다. (queryScore() 함수에서는 필요가 없어서 traversal2의 반환값을 저장하지 않는다)

색상 조합 확인 백트래킹 _ traversal2()

traversal2도 백트래킹의 원리를 이용한다. 다만 첫번째 백트래킹은 depth를 확인하기 위해 리프 노드에서 시작했다면 이번엔 서브트리의 색상 조합을 알기 위해 루트 노드에서 시작한다.

private static Set<Integer> traversal2(Node node) {
    // leaf 노드일 때
    Set<Integer> colorSet = new HashSet<>();
    if(node.getChildren().size() == 0) {
        result += 1;
        colorSet.add(node.getColor());
        return colorSet;
    }

    for(Node child : node.getChildren()) {
        colorSet.addAll(traversal2(child));
        // System.out.println(colorSet);
    }

    colorSet.add(node.getColor());
    result += (colorSet.size() * colorSet.size());
    return colorSet;
}

방법은 "후위 순회법"을 차용한다. 모든 서브트리의 값을 가져온 후에 본인의 노드를 확인한다. 리프 노드에 무슨 색상이 몇개 있는지 중요하지 않으므로 Set에 저장하였다.

전체 코드

import java.io.*;
import java.util.*;

public class Main {

    static class Node {
        int mId;
        int color;
        int maxDepth;
        Node parent = null;
        ArrayList<Node> children = new ArrayList<>();

        Node(int mId, int color, int maxDepth) {
            this.mId = mId;
            this.color = color;
            this.maxDepth = maxDepth;
        }

        void setParent(Node node) {
            this.parent = node;
        }

        void addChild(Node node) {
            this.children.add(node);
        }

        ArrayList<Node> getChildren() {
            return children;
        }

        void setColor(int color) {
            this.color = color;
        }

        int getColor() {
            return this.color;
        }
    }

    static ArrayList<Node> rootArr = new ArrayList<>();
    static Map<Integer, Node> nodeMap = new HashMap<>();
    static int result = 0;

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st;
        StringBuilder sb = new StringBuilder();

        int Q = Integer.parseInt(br.readLine());

        while(Q-- > 0) {
            st = new StringTokenizer(br.readLine());

            int command = Integer.parseInt(st.nextToken());
            int mId;
            int pId;
            int color;
            int maxDepth;
            switch (command) {
                case 100:
                    mId = Integer.parseInt(st.nextToken());
                    pId = Integer.parseInt(st.nextToken());
                    color = Integer.parseInt(st.nextToken());
                    maxDepth = Integer.parseInt(st.nextToken());
                    
                    addNode(new Node(mId, color, maxDepth), pId);
                    break;
                case 200:
                    mId = Integer.parseInt(st.nextToken());
                    color = Integer.parseInt(st.nextToken());
                    changeColor(mId, color);
                    break;
                case 300:
                    mId = Integer.parseInt(st.nextToken());
                    color = queryColor(mId);
                    sb.append(color).append("\n");
                    break;
                case 400:
                    sb.append(queryScore()).append("\n");
                    break;
            }
        }

        System.out.println(sb);
    }

    static void addNode(Node node, int pId) {
        if(pId == -1) {
            rootArr.add(node);
            nodeMap.put(node.mId, node);
            return;
        }

        // maxDepth 검증 백트래킹
        Node parent = nodeMap.get(pId);
        boolean enableFlag = check(parent, 2);
        if(enableFlag) {

            // 부모 자식 관계 설정
            node.setParent(parent);
            parent.addChild(node);

            // 빠른 검색을 위해 map에 저장
            nodeMap.put(node.mId, node);
        } 
    }

    static private boolean check(Node node, int depth) {
        // 루트일 때
        // System.out.println(node.maxDepth + " " + depth);
        if(node.parent == null) {
            // System.out.println("루트에서 --- " + node.maxDepth + " " + depth);
            if(node.maxDepth < depth) {
                return false;
            } else {
                return true;
            }
        }
        
        boolean enableFlag = check(node.parent, depth + 1);

        if(enableFlag) {
            if(node.maxDepth < depth) {
                return false;
            } else {
                return true;
            }
        } else {
            return false;
        }
    }

    static void changeColor(int mId, int color) {
        Node root = nodeMap.get(mId);
        root.setColor(color);
        // System.out.println("색 변경 ----" + root.mId + "번 노드의 색상은 " + color);

        // traversal
        traversal(root, color);
    }

    private static void traversal(Node node, int color) {

        for(Node child : node.getChildren()) {
            child.setColor(color);
            // System.out.println("색 변경 ----" + child.mId + "번 노드의 색상은 " + color);

            traversal(child, color);
        }
    }

    static int queryColor(int mId) {
        // System.out.println("색 조회 결과 ---- " + nodeMap.get(mId).getColor());

        return nodeMap.get(mId).getColor();
    }

    static int queryScore() {
        result = 0;

        for(Node root : rootArr) {
            traversal2(root);
        }
        // System.out.println("점수 조회 결과는 ---- " + result);

        return result;
    }

    private static Set<Integer> traversal2(Node node) {
        // leaf 노드일 때
        Set<Integer> colorSet = new HashSet<>();
        if(node.getChildren().size() == 0) {
            result += 1;
            colorSet.add(node.getColor());
            return colorSet;
        }

        for(Node child : node.getChildren()) {
            colorSet.addAll(traversal2(child));
            // System.out.println(colorSet);
        }

        colorSet.add(node.getColor());
        result += (colorSet.size() * colorSet.size());
        return colorSet;
    }
}

풀이 소요 시간 : 1시간 30분

트리를 이해하고 있다면 크게 어려운 문제가 아니다. 문제에 주어진대로 차분히 구현하면 더 빠른 시간안에 풀 수 있을 것이다.