[해커랭크] Contacts [JAVA]

문제링크 : www.hackerrank.com/challenges/contacts/problem

Contacts | HackerRank

정말 오랜만에 트라이 문제를 풀었다.

public static Node root;

    public static class Node{
        char ch;
        boolean is_end;
        int cnt;
        Node child[] = new Node[26];
        
    }
    public static void add(String str){
        int idx = 0;
        Node cur = root;
        while(idx!=str.length()){
            int chNum = str.charAt(idx)-'a';
            if(cur.child[chNum]==null){
                cur.child[chNum] = new Node();
            }
            idx++;
            cur = cur.child[chNum];
            cur.cnt++;
        }

        cur.is_end = true;
    }
    public static int find(String str){
        Node cur = root;
        int idx = 0;
        while(idx!=str.length()){
            int chNum = str.charAt(idx)-'a';
            if(cur.child[chNum]==null) return 0;
            idx++;
            cur = cur.child[chNum];
        }
        return cur.cnt;
    }
    static int[] contacts(String[][] queries) {
        root = new Node();
        ArrayList<Integer> ans = new ArrayList<Integer>();
        for(int i=0;i<queries.length;i++){
            if(queries[i][0].equals("add")) add(queries[i][1]);
            else{
                ans.add(find(queries[i][1]));
            }
        }
        int ret[] = new int[ans.size()];
        for(int i=0;i<ans.size();i++){
            ret[i] = ans.get(i);
        }
        return ret;
    }