692. Top K Frequent Words

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
class Solution {
private static class Word implements Comparable<Word> {
private final String word;
private final int freq;

public Word(String word, int freq) {
this.word = word;
this.freq = freq;
}

@Override
public int compareTo(Word o) {
int diff = o.freq - this.freq;
if (diff == 0) {
return this.word.compareTo(o.word);
} else {
return diff;
}
}
}

public List<String> topKFrequent(String[] words, int k) {
Map<String, Integer> wordToFreqMap = new HashMap<>();
for (String word : words) {
wordToFreqMap.put(word, wordToFreqMap.getOrDefault(word, 0) + 1);
}

Word[] wordArray = wordToFreqMap.entrySet().stream().map(entry -> new Word(entry.getKey(), entry.getValue())).toArray(Word[]::new);
int targetIndex = k - 1;
int left = 0, right = wordArray.length - 1;
while (true) {
int pivotIndex = partition(wordArray, left, right);
if (pivotIndex < targetIndex) {
left = pivotIndex + 1;
} else if (pivotIndex > targetIndex) {
right = pivotIndex - 1;
} else {
Word[] targetArray = Arrays.copyOfRange(wordArray, 0, pivotIndex + 1);
Arrays.sort(targetArray);
return Arrays.stream(targetArray).map(word -> word.word).collect(Collectors.toList());
}
}
}

private int partition(Word[] wordArray, int left, int right) {
int randomIndex = left + ThreadLocalRandom.current().nextInt(right - left + 1);
swap(wordArray, left, randomIndex);

Word pivotWord = wordArray[left];
int i = left, j = right;
while (i < j) {
while (i < j && wordArray[j].compareTo(pivotWord) >= 0) {
j--;
}
while (i < j && wordArray[i].compareTo(pivotWord) <= 0) {
i++;
}
swap(wordArray, i, j);
}
swap(wordArray, i, left);
return i;
}

private void swap(Word[] wordArray, int i, int j) {
Word tmp = wordArray[i];
wordArray[i] = wordArray[j];
wordArray[j] = tmp;
}
}

References

692. Top K Frequent Words