1407. Problem - Word Rank
Prefix


Implement data structure for searching word based on rank.

1. Requirement

Design a data structure to store the given words and search words with prefix.

2. Solution

Word bean.

public class Word {
    public String name;
    public int rank;
    public Word(String name, int rank) {
        this.name = name;
        this.rank = rank;
    }
}

Trie.

public class Trie {
    class TrieNode {
        public Map<Character, TrieNode> children;
        public int rank; // only valid when leaf = true
        public boolean leaf;

        public TrieNode() {
            children = new HashMap<>();
            leaf = false;
        }
    }

    private TrieNode root;

    public Trie() {
        this.root = new TrieNode();
    }

    public TrieNode getRoot() {
        return this.root;
    }

    // Return true if the word is in trie
    public boolean search(String word) {
        TrieNode tn = searchNode(word);
        if (tn != null && tn.leaf) {
            return true;
        } else {
            return false;
        }
    }

    // Return true if there is any word in trie that starts with the given prefix
    public boolean startsWith(String prefix) {
        if (searchNode(prefix) == null) {
            return false;
        } else {
            return true;
        }
    }

    // Return true if there is any word in trie that starts with the given prefix
    public List<Word> searchWords(String prefix) {
        TrieNode current = root;
        StringBuilder sb = new StringBuilder();
        List<Word> list = new ArrayList<>();

        for (int i = 0; i < prefix.length(); i++) {
            char ch = prefix.charAt(i);
            if (!current.children.containsKey(ch)) {
                return null;
            } else {
                sb.append(ch);
                current = current.children.get(ch);
            }
        }

        dfs(current, sb.toString(), list);

        return list;
    }

    private void dfs(TrieNode node, String prefix, List<Word> list) {
        if (node.leaf) {
            list.add(new Word(prefix, node.rank));
        }

        for (Map.Entry<Character, TrieNode> entry : node.children.entrySet()) {
            dfs(entry.getValue(), prefix + entry.getKey(), list);
        }
    }

    private TrieNode searchNode(String str) {
        TrieNode current = root;

        for (int i = 0; i < str.length(); i++) {
            char ch = str.charAt(i);
            if (current.children.containsKey(ch)) {
                current = current.children.get(ch);
            } else {
                return null;
            }
        }

        return current;
    }

    // Insert a word into trie
    public void insert(String word, int rank) {
        TrieNode current = root;

        for (int i = 0; i < word.length(); i++) {
            char ch = word.charAt(i);
            if (!current.children.containsKey(ch)) {
                current.children.put(ch, new TrieNode());
            }
            current = current.children.get(ch);
        }

        current.rank = rank;
        current.leaf = true;
    }

    public boolean delete(String word) {
        TrieNode current = root;
        TrieNode lastBranchNode = null;
        Character lastBrachChar = null;

        for (int i = 0; i < word.length(); i++) {
            char ch = word.charAt(i);
            if (current.children.containsKey(ch)) {
                if (current.children.size() > 1) {
                    lastBranchNode = current;
                    lastBrachChar = ch;
                }
                current = current.children.get(ch);
            } else {
                // word not found
                return false;
            }
        }

        if (current.children.size() > 0) {
            // case 1: The to-be deleted word is prefix of another long word in trie.
            current.leaf = false;
            return true;
        }

        if (lastBranchNode != null) {
            // case 2: The to-be deleted word has other words as prefix
            lastBranchNode.children.remove(lastBrachChar);
            return true;
        } else {
            // case 3: The to-be deleted word present as unique word
            root.children.remove(word.charAt(0));
            return true;
        }
    }
}

WorkRank.

public class WordRank {
    private Trie trie;

    public WordRank(List<Word> words) {
        trie = new Trie();
        for (Word word : words) {
            trie.insert(word.name, word.rank);
        }
    }

    public List<Word> search(String prefix) {
        return trie.searchWords(prefix).stream().sorted((w1, w2) -> w1.rank - w2.rank).collect(Collectors.toList());
    }
}

The example how to use it.

public class WordRankExample {
    private static final String INPUT_FILE = "input.txt";
    private static final String PREFIX_FILE = "prefix.txt";
    private static final String OUTPUT_FILE = "output.txt";

    public static void main(String args[]) throws Exception {

        ClassLoader classLoader = WordRankExample.class.getClassLoader();

        // Get words from file
        List<Word> words = new ArrayList<>();
        Path path = Paths.get("files", INPUT_FILE);
        try (InputStream inputStream = classLoader.getResourceAsStream(path.toString())) {
            System.setIn(inputStream);
            Scanner sc = new Scanner(System.in);
            while (sc.hasNextLine()) {
                words.add(new Word(sc.next(), sc.nextInt()));
            }
        } catch (IOException e) {
            e.printStackTrace();
        }

        // Create Work Rank object
        WordRank wr = new WordRank(words);

        // Get prefixes from file
        List<String> prefixes = new ArrayList<>();
        path = Paths.get("files", PREFIX_FILE);
        try (InputStream inputStream = classLoader.getResourceAsStream(path.toString())) {
            System.setIn(inputStream);
            Scanner sc = new Scanner(System.in);
            while (sc.hasNext()) {
                prefixes.add(sc.next());
            }
        } catch (IOException e) {
            e.printStackTrace();
        }

        // Search
        for (String pre : prefixes) {
            System.out.println(pre + ":");
            List<Word> list = wr.search(pre);
            for (Word word : list) {
                System.out.println(word.name + " " + word.rank);
            }
            System.out.println();
        }
    }
}

Input.

hello 6
world 10
wide 3
hell 4
worldwide 7
lyft 20

Prefix.

hell
world

Output.

hell:
hell 4
hello 6

world:
worldwide 7
world 10

Test class.

public class WordRankTest {
    @Test
    public void testWordRank() {
        System.out.println("testWordRank");

        List<Word> words = new ArrayList<>();
        words.add(new Word("hello", 6));
        words.add(new Word("world", 10));
        words.add(new Word("wide", 3));
        words.add(new Word("hell", 4));
        words.add(new Word("worldwide", 7));
        words.add(new Word("lyft", 20));

        WordRank wr = new WordRank(words);
        List<Word> res1 = wr.search("hell");
        assertEquals(2, res1.size());
        assertEquals("hell", res1.get(0).name);
        assertEquals(4, res1.get(0).rank);
        assertEquals("hello", res1.get(1).name);
        assertEquals(6, res1.get(1).rank);

        List<Word> res2 = wr.search("world");
        assertEquals(2, res2.size());
        assertEquals("worldwide", res2.get(0).name);
        assertEquals(7, res2.get(0).rank);
        assertEquals("world", res2.get(1).name);
        assertEquals(10, res2.get(1).rank);
    }
}

3. Source Files