[알고리즘, 코드트리] 코드트리 메신저 - java

 

  ※ 주의!! 이 문제는 사람의 인내심과 정신력을 시험합니다!! 어쭙잖은 각오로 이 문제를 도전하지 마세요.  

문제

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

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

요약하자면 아래와 같다. 총 다섯 가지 동작을 수행해야 한다.

(1) 사내 메신저 준비 : 이진트리 구조의 사내 메신저를 초기화합니다. 

(2) 알림망 설정 ON / OFF  : 처음 모든 채팅방의 알림망 설정은 켜져있습니다. 이 기능이 작동되면 c번 채팅방의 알림망 설정이 ON 상태라면 OFF로 바꿔주고, OFF 상태라면 ON으로 바꿔줍니다. 알림망 설정이 OFF 상태가 되면 자기 자신을 포함하여 아래에서 올라온 모든 알림을 더 이상 위로 올려 보내지 않습니다.

(3) 권한 세기 변경 : c번 채팅방의 권한 세기를 power로 변경합니다.

(4) 부모 채팅방 교환 : c1​번 채팅방과 c2​번 채팅방의 부모를 서로 바꿉니다. 이때 c1​번 채팅방과 c2​번 채팅방은 트리 내에서 같은 depth상에 있음을 가정해도 좋습니다.

(5) 알림을 받을 수 있는 채팅방 수 조회 : 메세지를 보냈을 때 c번 채팅방까지 알림이 도달할 수 있는 서로 다른 채팅방의 수를 출력

입력

첫 번째 줄에 채팅방의 수 N과 명령의 수 Q가 공백을 사이에 두고 주어집니다.

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

• 사내 메신저 준비
• 100 p1 p2 ... pN a1 a2 .... aN 형태로 공백을 사이에 두고 주어집니다. 1번부터 N번까지 각 채팅방의 부모 채팅방 번호가 순서대로 p1​부터 pN​까지고, 초기 권한 세기가 각각 a1​부터 an​까지임을 뜻합니다. 이 명령은 항상 첫 번째 명령으로만 주어지며, 항상 주어집니다. 0번 채팅방의 p값과 a값은 주어지지 않음에 유의합니다.
• 알림망 설정 ON / OFF
• 200 c 형태로 공백을 사이에 두고 주어집니다. 이는 c번 채팅방의 알림망 설정이 ON 상태라면 OFF로 바꿔주고, OFF 상태라면 ON 상태로 바꿔줘야 함을 의미합니다.
• 권한 세기 변경
• 300 c power 형태로 공백을 사이에 두고 주어집니다. 이는 c번 채팅방의 권한 세기를 power로 변경해야 함을 의미합니다.
• 부모 채팅방 교환
• 400 c1 c2 형태로 공백을 사이에 두고 주어집니다. 이는 c1​번 채팅방과 c2​번 채팅방의 부모를 변경해야 함을 의미합니다. 이때 c1​번 채팅방과 c2​번 채팅방은 트리 내에서 같은 depth상에 있음을 가정해도 좋습니다.
• 알림을 받을 수 있는 채팅방 수 조회
• 500 c 형태로 공백을 사이에 두고 주어집니다. 이 경우 메세지를 보냈을 때 c번 채팅방까지 알림이 도달할 수 있는 서로 다른 채팅방의 수를 출력합니다. 단, 이때 c번 채팅방을 제외한 채팅방의 수를 세어야 함에 유의합니다.

• 1 ≤ Q ≤ 100,000
• 1 ≤ N ≤ 100,000
• 1 ≤ 주어지는 이진트리의 최대 깊이(D) ≤ 20
• 0 ≤ pi​ ≤ N
• 1 ≤ ai​,c,power,c1​,c2​ ≤ N
• c1​ != c2​
  

출력

500 명령어에 대해 c번 채팅방까지 알림이 도달할 수 있는 서로 다른 채팅방의 수를 한 줄에 하나씩 출력합니다.

문제 풀이

시간 제한의 늪에 빠져서 5시간을 허우적거린 문제였다. 결국 푸는데 성공했지만... 5시간 걸려서 풀었으면 성공이라고 얘기할 수 있나..? 왜 시간 제한에 걸렸는지도 해설이 끝나고 설명해보겠다. 풀이를 시작해보자.

자료구조 정의 및 입출력

static class Node {
    int id;
    int authority;
    boolean alram = true;

    int parentId;
    Node parent;

    // 현재 노드가 전달할 수 있는 알람 수
    // arrOfAlram[3] = 4 라면 위로 3개의 노드에 전달할 수 있는 알람이 4
    int[] arrOfAlram = new int[21];

    Node rightChild;
    int rightMaxAuthority;

    Node leftChild;
    int leftMaxAuthority;

    Node(int id, int authority, int parentId) {
        this.id = id;
        this.authority = authority;
        this.parentId = parentId;
    }
}

static Map<Integer, Node> nodeMap = new HashMap<>();

문제에서도 친절히 설명한대로 자료 구조는 이진 트리이다. 따라서 Node 클래스를 선언하고  rightChild와 leftChild를 받아 이진 트리를 구성하자. 트리의 구현을 연습했다면 이정도 자료구조를 생각하는 것은 상당히 간단할 것이다. 

 

이 문제의 핵심은 현재 노드가 상위 노드로 전달할 수 있는 알람 수를 저장하는 것이다. 이유는 후술한다. 따라서 Node 클래스의 멤버 변수로 트리의 최대 뎁스(20) 크기의 arrOfAlram이라는 int형 배열을 선언한다.

N = Integer.parseInt(st.nextToken());
Q = Integer.parseInt(st.nextToken());

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

    int command = Integer.parseInt(st.nextToken());

    switch (command) {
        case 100: 
            init(st);
            break;
        case 200 :
            c1 = Integer.parseInt(st.nextToken());
            turn_alram(c1);
            break;
        case 300 :
            c1 = Integer.parseInt(st.nextToken());
            power = Integer.parseInt(st.nextToken());
            change_authority(c1, power);
            break;
        case 400 :
            c1 = Integer.parseInt(st.nextToken());
            c2 = Integer.parseInt(st.nextToken());
            change_parent(c1, c2);
            break;
        case 500: 
            c1 = Integer.parseInt(st.nextToken());
            int count = getArlam(c1);
            sb.append(count).append("\n");
    }
}
System.out.println(sb);

각 명령(command)에 대해서 동작을 수행할 함수를 작성한다. 구조를 다 잡았으니 하나하나 구현해보자.

command 100 : 이진 트리 초기화 _ init()

static private void init(StringTokenizer st) {
    // root 생성
    nodeMap.put(0, new Node(0, 0, -1));
    Node root = nodeMap.get(0);

    // 다른 노드 생성만
    for(int i= 1; i <= N; i++) {
        nodeMap.put(i, new Node(i, 0, -1));
    }

    // 노드 연결
    for(int i = 1; i <= N; i++) {
        int parentId = Integer.parseInt(st.nextToken());

        Node parent = nodeMap.get(parentId);
        Node child = nodeMap.get(i);

        if(parent.leftChild == null) {
            parent.leftChild = child;
            child.parent = parent;
            child.parentId = parentId;
        } else {
            parent.rightChild = child;
            child.parent = parent;
            child.parentId = parentId;
        }
    }

    // 권한 설정
    for(int i = 1; i <= N; i++) {
        int authority = Integer.parseInt(st.nextToken());

        Node node = nodeMap.get(i);
        node.authority = authority;
    }

    // ArrOfAlram 초기화
    alramInit(root);
}

단순히 입력 받은 값에 이진트리를 만드는 과정이므로 어렵지 않다. 특별히 설명할 것은 없으므로 arrOfAlram 배열을 초기화하는 함수만 언급하고 넘어가겠다.

arrOfAlram 배열 초기화 _ alramInit()

static private void alramInit(Node now) {
    if(now == null) {
        return;
    }

    now.arrOfAlram = new int[21];
    for(int i = 0; i <= now.authority && i <= 20; i++) {
        now.arrOfAlram[i]++;
    }

    alramInit(now.leftChild);
    alramInit(now.rightChild);

    if(now.leftChild != null) {
        // lefChilde의 n(1~20)번 알람 수가 현재 노드의 n - 1번 알람 수값에 반영
        for(int i = 1; i <= 20; i++) {
            now.arrOfAlram[i - 1] += now.leftChild.arrOfAlram[i];
        }
    }

    if(now.rightChild != null) {
        for(int i = 1; i <= 20; i++) {
            now.arrOfAlram[i - 1] += now.rightChild.arrOfAlram[i];
        }
    }
}

혹시라도 파라미터로 넘어온 node가 null일 수 있으므로 가장 먼저 Null 체크를 해준다.

 

그 후 본인의 권한으로 상위 노드로 올려보낼 수 있는 만큼 배열에 1을 더해준다.

위 배열에서 id:8 노드의 arrOfAlram[2]의 값은 8노드 기준 두번째 상위 노드까지 본인의 알람을 전달할 수 있다는 의미가 된다.

 

id:4 노드 기준으로 다시 생각해보자.

본인의 권한에 의해서 4의 arrOfAlram은 위와같이 초기화된다. 이 때, 4는 7과 8에서 알람을 받게 될 것이다. 다음 과정을 반드시 이해하고 넘어가자.

4는 본인의 자식 노드인 7과 8번 노드에서 알람을 넘겨 받게 된다. 따라서, 4번 노드는 본인의 알람을 포함하여 총 3개의 알람이 울리고(arrOfAlram[0] == 3), 본인의 바로 한 단계 상위 노드인 1번 노드로는 2개의 알람을 전송하게 된다(arrOfAlram[1] == 2).

 

마지막으로 1번 노드를 확인하고 원리를 마무리하자.

1의 권한은 1이므로 본인의 arrOfAlram 배열을 다음과 같이 초기화한다.

이제 자식 노드로부터 알람의 수를 넘겨받자.

따라서 1번 노드는 본인의 알람을 포함하여 4개의 알람을 받고(arrOfAlram[0] == 4) 바로 상위 노드인 루트 노드로 3개의 알람을 보내게 된다(arrOfAlram[1] == 3).

command 500 : 받을 수 있는 알람 수는?

Node에 이렇게 상태를 저장하고 있는 arrOfAlram 배열이 있다면 500번 명령어의 연산은 단순히 O(1)로 끝날 수 있다.

static private int getArlam(int c) {
    Node node = nodeMap.get(c);

    return node.arrOfAlram[0] - 1;
}

본인의 알람은 답에서 제외하므로 1을 빼고 리턴하자.

command 200 : 알람 활성화/비활성화 _ turn_alram()

static private void turn_alram(int c) {
    Node node = nodeMap.get(c);
    node.alram = !node.alram;

    // node부터 root까지 거슬러 올라가면서 arrOfAlram 재계산
    recal(c);
}

특정 노드에 대해서 알람망 기능을 끄게 된다면 해당 노드를 루트로 하는 서브 트리의 알람은 기존의 이진 트리에 영향을 끼칠 수 없다. "단, 분리된 서브 트리 내에서는 새로운 알람 네트워크가 형성된 것에 유의하자."

이를 반영하여 arrOfAlram을 재계산할 수 있도록 새로운 함수를 작성하자.

변경된 상황을 반영하는 함수 _ recal()

static private void recal(int id) {
    Node now = nodeMap.get(id);

    while(now != null && now.id != 0) {
        Arrays.fill(now.arrOfAlram, 0);
        for(int i = 0; i <= now.authority && i <= 20; i++) {
            now.arrOfAlram[i]++;
        }

        if(now.leftChild != null && now.leftChild.alram) {
            // leftChild의 n(1~20)번 알람 수가 현재 노드의 n - 1번 알람 수값에 반영
            for(int i = 1; i <= 20 ; i++) {
                now.arrOfAlram[i - 1] += now.leftChild.arrOfAlram[i];
            }
        }

        if(now.rightChild != null && now.rightChild.alram) {
            for(int i = 1; i <= 20; i++) {
                now.arrOfAlram[i - 1] += now.rightChild.arrOfAlram[i];
            }
        }

        now = now.parent;
    }
}

변경이 생긴 노드부터 재계산을 시작하여 상위 노드로 거슬러 올라간다. 연산은 노드가 루트에 도달했을 때 마무리한다(혹시 몰라서 널 체크도 추가했다).

단, 재계산 함수에서는 자식 노드의 알람 기능을 확인해야 한다. 자식 노드의 알람 기능이 켜져있을 때(alram == true)만 현재 노드로 결과를 반영하도록 설계하자.

command 300 : 권한 세기 변경 _ change_authority()

static private void change_authority(int c, int power) {
    Node node = nodeMap.get(c);
    node.authority = power;

    recal(c);
}

recal함수를 정의해 두어서 300 명령어에 대해서는 크게 해줄 일이 없다.

command 400 : 부모 변경 _ change_parent()

static private void change_parent(int c1, int c2) {
    Node node1 = nodeMap.get(c1);
    Node node2 = nodeMap.get(c2);

    Node parent1 = nodeMap.get(node1.parentId);
    Node parent2 = nodeMap.get(node2.parentId);

    if(parent1.leftChild == node1) {
        parent1.leftChild = node2;
    } else if(parent1.rightChild == node1) {
        parent1.rightChild = node2;
    } 
    node2.parentId = parent1.id;
    node2.parent = parent1;

    if(parent2.leftChild == node2) {
        parent2.leftChild = node1;

    } else if(parent2.rightChild == node2) {
        parent2.rightChild = node1;
    }
    node1.parentId = parent2.id;
    node1.parent = parent2;

    recal(node1.id);
    recal(node2.id);
}

연결리스트 swap과 유사한 구조의 코드이다. 이 또한, 이 문제에 도전할 정도의 수리논술력이면 이해하기 어려울 거라고 생각하지 않는다.

이 경우도 recal 함수를 호출하여 변경된 상황을 다시 반영한다.

전체 코드

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

public class Main {
    static class Node {
        int id;
        int authority;
        boolean alram = true;

        int parentId;
        Node parent;

        // 현재 노드가 전달할 수 있는 알람 수
        // arrOfArlam[3] = 4 라면 위로 3개의 노드에 전달할 수 있는 알람이 4
        int[] arrOfAlram = new int[21];

        Node rightChild;
        int rightMaxAuthority;

        Node leftChild;
        int leftMaxAuthority;

        Node(int id, int authority, int parentId) {
            this.id = id;
            this.authority = authority;
            this.parentId = parentId;
        }
    }

    static Map<Integer, Node> nodeMap = new HashMap<>();

    static int N, Q;

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

        st = new StringTokenizer(br.readLine());

        N = Integer.parseInt(st.nextToken());
        Q = Integer.parseInt(st.nextToken());

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

            int command = Integer.parseInt(st.nextToken());

            switch (command) {
                case 100: 
                    init(st);
                    break;
                case 200 :
                    c1 = Integer.parseInt(st.nextToken());
                    turn_alram(c1);
                    break;
                case 300 :
                    c1 = Integer.parseInt(st.nextToken());
                    power = Integer.parseInt(st.nextToken());
                    change_authority(c1, power);
                    break;
                case 400 :
                    c1 = Integer.parseInt(st.nextToken());
                    c2 = Integer.parseInt(st.nextToken());
                    change_parent(c1, c2);
                    break;
                case 500: 
                    c1 = Integer.parseInt(st.nextToken());
                    int count = getArlam(c1);
                    sb.append(count).append("\n");
            }
        }
        System.out.println(sb);
    }

    static private void init(StringTokenizer st) {
        // root 생성
        nodeMap.put(0, new Node(0, 0, -1));
        Node root = nodeMap.get(0);

        // 다른 노드 생성만
        for(int i= 1; i <= N; i++) {
            nodeMap.put(i, new Node(i, 0, -1));
        }

        // 노드 연결
        for(int i = 1; i <= N; i++) {
            int parentId = Integer.parseInt(st.nextToken());

            Node parent = nodeMap.get(parentId);
            Node child = nodeMap.get(i);

            if(parent.leftChild == null) {
                parent.leftChild = child;
                child.parent = parent;
                child.parentId = parentId;
            } else {
                parent.rightChild = child;
                child.parent = parent;
                child.parentId = parentId;
            }
        }

        // 권한 설정
        for(int i = 1; i <= N; i++) {
            int authority = Integer.parseInt(st.nextToken());

            Node node = nodeMap.get(i);
            node.authority = authority;
        }

        // ArrOfAlram 초기화
        alramInit(root);
    }
    
    static private void alramInit(Node now) {
        if(now == null) {
            return;
        }

        now.arrOfAlram = new int[21];
        for(int i = 0; i <= now.authority && i <= 20; i++) {
            now.arrOfAlram[i]++;
        }
        
        alramInit(now.leftChild);
        alramInit(now.rightChild);
        
        if(now.leftChild != null) {
            // lefChilde의 n(1~20)번 알람 수가 현재 노드의 n - 1번 알람 수값에 반영
            for(int i = 1; i <= 20; i++) {
                now.arrOfAlram[i - 1] += now.leftChild.arrOfAlram[i];
            }
        }

        if(now.rightChild != null) {
            for(int i = 1; i <= 20; i++) {
                now.arrOfAlram[i - 1] += now.rightChild.arrOfAlram[i];
            }
        }
    }

    static private void recal(int id) {
        Node now = nodeMap.get(id);

        while(now != null && now.id != 0) {
            Arrays.fill(now.arrOfAlram, 0);
            for(int i = 0; i <= now.authority && i <= 20; i++) {
                now.arrOfAlram[i]++;
            }

            if(now.leftChild != null && now.leftChild.alram) {
            // leftChild의 n(1~20)번 알람 수가 현재 노드의 n - 1번 알람 수값에 반영
                for(int i = 1; i <= 20 ; i++) {
                    now.arrOfAlram[i - 1] += now.leftChild.arrOfAlram[i];
                }
            }

            if(now.rightChild != null && now.rightChild.alram) {
                for(int i = 1; i <= 20; i++) {
                    now.arrOfAlram[i - 1] += now.rightChild.arrOfAlram[i];
                }
            }

            now = now.parent;
        }
    }

    static private void turn_alram(int c) {
        Node node = nodeMap.get(c);
        node.alram = !node.alram;

        // node부터 root까지 거슬러 올라가면서 arrOfAlram 재계산
        recal(c);
    }

    static private void change_authority(int c, int power) {
        Node node = nodeMap.get(c);
        node.authority = power;

        recal(c);
    }

    static private void change_parent(int c1, int c2) {
        Node node1 = nodeMap.get(c1);
        Node node2 = nodeMap.get(c2);

        Node parent1 = nodeMap.get(node1.parentId);
        Node parent2 = nodeMap.get(node2.parentId);

        if(parent1.leftChild == node1) {
            parent1.leftChild = node2;
        } else if(parent1.rightChild == node1) {
            parent1.rightChild = node2;
        } 
        node2.parentId = parent1.id;
        node2.parent = parent1;

        if(parent2.leftChild == node2) {
            parent2.leftChild = node1;
            
        } else if(parent2.rightChild == node2) {
            parent2.rightChild = node1;
        }
        node1.parentId = parent2.id;
        node1.parent = parent2;

        recal(node1.id);
        recal(node2.id);
    }

    static private int getArlam(int c) {
        Node node = nodeMap.get(c);

        return node.arrOfAlram[0] - 1;
    }
}

문제 풀이 소요 시간 : 5시간 이상

 

상당히 난이도가 있는 문제였다. 몇가지 더 언급하고 마무리해보자. 왜 내 원래 코드는 시간 초과가 났을까? 기존의 코드에서는 500번 명령어에서 트리의 순회로 직접 depth와 authority를 비교하는 계산을 수행하였다. 자, 인풋 사이즈를 보자.

트리의 순회로 모든 노드를 다 탐색한다고 가정하면 500번 명령어의 시간 복잡도는 O(N)이다. 만약 최대 인풋 사이즈에서 O(N) 함수를 Q의 최대인 10만 번 호출한다고 가정하면 전체 함수의 시간복잡도는 O(N x Q)로 100억번의 연산을 수행해야 한다. 시간 초과는 당연했다.

500번 명령어의 시간복잡도를 O(1)로 낮추는 것이 이 문제의 핵심이다.

 

따라서 이진트리의 최대 깊이가 20이라는 점을 이용하기로 생각했더니 풀이가 보였다. 각 노드마다 이진트리의 최대 깊이만큼의 배열을 가지고 여기에 정보를 저장한다면 변경이 생겨 매번 재계산 해줘도 O(D x Q)로 계산수를 획기적으로 줄일 수 있다.