Distance Code

//thinking depth first with iterative deepening maybe (though slow)
//using heap

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.HashSet;
import java.util.PriorityQueue;

public class Distance_Code_v3 {

    private static class Node implements Comparable<Node>{

        int number;
        HashSet<Node> nbs = new HashSet<>();

        public Node(int number){
            this.number = number;
        }

        @Override
        public int compareTo(Node o) {
            return Integer.compare(this.nbs.size(), o.nbs.size());
        }
    }

    private static class Node2{

        int number;
        ArrayList<Node2> children = new ArrayList<>();

        public Node2(int number){
            this.number = number;
        }
    }

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        int type = Integer.valueOf(br.readLine());
        int N = Integer.valueOf(br.readLine());
        if(type == 1){
            ArrayList<Node> nodes = new ArrayList<>();
            for(int i = 0; i < N; i++) nodes.add(null);
            for(int i = 0; i < N - 1; i++){
                String[] input = br.readLine().split(" ");
                int a = Integer.valueOf(input[0])-1;
                int b = Integer.valueOf(input[1])-1;
                if(nodes.get(a) == null){
                    nodes.set(a, new Node(a));
                }
                if(nodes.get(b) == null) {
                    nodes.set(b, new Node(b));
                }
                nodes.get(a).nbs.add(nodes.get(b));
                nodes.get(b).nbs.add(nodes.get(a));
            }
            nodes.sort(new Comparator<Node>() {
                @Override
                public int compare(Node o1, Node o2) {
                    return Integer.compare(o1.nbs.size(), o2.nbs.size());
                }
            });
            PriorityQueue<Node> heap = new PriorityQueue<>(nodes);
            ArrayList<Integer> result = new ArrayList<>();
            while(!heap.isEmpty()){
                Node curr = heap.poll();
                result.add(curr.number);
                curr.nbs.remove(curr);
            }
            for(int i = 0; i < result.size(); i++){
                System.out.print((result.get(i)+1) + " ");
            }
        }else{
            inputs = br.readLine().split(" ");
            //if distance == 2 -> same depth && same parent
            //if distance == 1 -> its the parent
            //if distance == 3 -> different depth "different parent"
            //if distance == 4 -> same depth different parent
            //no its wrong
            inputIndex = inputs.length-1;
            Node2 root = new Node2(nextNumber);
            ArrayList<Node2> temp = new ArrayList<>();
            temp.add(root);
            decode(temp, 1);
            dfs(root);
        }
    }

    private static int nextNumber = 1;
    private static int inputIndex;
    private static String[] inputs;

    private static void decode(ArrayList<Node2> siblings, int depth){
        int siblingIndex = 0;
        while(inputIndex > -1){
            int diff = Integer.valueOf(inputs[inputIndex]);
            inputIndex--;
            if(diff == 1 || diff == 2 || diff == depth+1){
                nextNumber++;
                Node2 child = new Node2(nextNumber);
                siblings.get(siblingIndex).children.add(child);
            }else if(diff == 2*depth){ //assuming a depth shift is never on the curr sibling
                siblingIndex++;
                nextNumber++;
                Node2 child = new Node2(nextNumber);
                siblings.get(siblingIndex).children.add(child);
            }else{
                inputIndex++;
                break;
            }
        }
        for(int i = 0; i < siblings.size(); i++){
            if(siblings.get(i).children.size() != 0){
                decode(siblings.get(i).children, depth + 1);
            }
        }
    }

    private static void dfs(Node2 node){
        for(int i = 0; i < node.children.size(); i++){
            System.out.println(node.number + " " + node.children.get(i).number);
            dfs(node.children.get(i));
        }
    }
}

Test 1

Group: 1, 2, 3

Verdict: ACCEPTED

Test 2

Group: 1, 2, 3

Verdict: ACCEPTED

Test 3

Group: 1, 2, 3

Verdict: ACCEPTED

Test 4

Group: 1, 2, 3

Verdict: ACCEPTED

Test 5

Group: 1, 2, 3

Verdict: WRONG ANSWER

Test 6

Group: 1, 2, 3

Verdict: ACCEPTED

Test 7

Group: 1, 2, 3

Verdict: ACCEPTED

Test 8

Group: 1, 2, 3

Verdict: ACCEPTED

Test 9

Group: 1, 2, 3

Verdict: WRONG ANSWER

Test 10

Group: 1, 2, 3

Verdict: WRONG ANSWER

Test 11

Group: 1, 2, 3

Verdict: WRONG ANSWER

Test 12

Group: 2, 3

Verdict: ACCEPTED

Test 13

Group: 2, 3

Verdict: WRONG ANSWER

Test 14

Group: 2, 3

Verdict: WRONG ANSWER

Test 15

Group: 2, 3

Verdict: WRONG ANSWER

Test 16

Group: 3

Verdict: TIME LIMIT EXCEEDED

Test 17

Group: 3

Verdict: TIME LIMIT EXCEEDED

Test 18

Group: 3

Verdict: TIME LIMIT EXCEEDED

Test 19

Group: 3

Verdict: TIME LIMIT EXCEEDED

Test 20

Group: 1, 2, 3

Verdict: ACCEPTED