1+ use std:: collections:: HashMap ;
2+
3+ struct AutocompleteSystem {
4+ content : Vec < u8 > ,
5+ trie : Trie ,
6+ hash : HashMap < String , i32 > ,
7+ empty : bool ,
8+ }
9+
10+ impl AutocompleteSystem {
11+
12+ fn new ( sentences : Vec < String > , times : Vec < i32 > ) -> Self {
13+ let mut hash = HashMap :: new ( ) ;
14+ let mut trie = Trie :: new ( ) ;
15+ for ( s, t) in sentences. into_iter ( ) . zip ( times) {
16+ trie. insert ( s. as_bytes ( ) ) ;
17+ hash. insert ( s, t) ;
18+ }
19+ Self {
20+ content : vec ! [ ] ,
21+ trie,
22+ hash,
23+ empty : false ,
24+ }
25+ }
26+
27+ fn input ( & mut self , c : char ) -> Vec < String > {
28+ let mut ans = vec ! [ ] ;
29+ if c == '#' {
30+ * self . hash . entry ( unsafe { std:: str:: from_utf8_unchecked ( & self . content ) . to_string ( ) } ) . or_insert ( 0 ) += 1 ;
31+ self . empty = false ;
32+ self . trie . insert ( & self . content ) ;
33+ self . content . clear ( ) ;
34+ return vec ! [ ] ;
35+ } else {
36+ self . content . push ( c as u8 ) ;
37+ if self . empty {
38+ return vec ! [ ] ;
39+ }
40+ ans = self . trie . find_prefix ( & self . content , 0 )
41+ }
42+ if ans. is_empty ( ) {
43+ self . empty = true ;
44+ }
45+ if self . empty {
46+ return vec ! [ ] ;
47+ }
48+ ans. sort_by ( |a, b| {
49+ let ta = * self . hash . get ( a) . unwrap ( ) ;
50+ let tb = * self . hash . get ( b) . unwrap ( ) ;
51+ if ta == tb {
52+ a. cmp ( b)
53+ } else {
54+ tb. cmp ( & ta)
55+ }
56+ } ) ;
57+
58+ ( 0 ..ans. len ( ) . min ( 3 ) ) . into_iter ( ) . map ( |i| ans[ i] . clone ( ) ) . collect ( )
59+ }
60+ }
61+
62+ struct Trie {
63+ data : HashMap < u8 , Option < Box < Trie > > > ,
64+ term : bool ,
65+ }
66+
67+ impl Trie {
68+ fn new ( ) -> Self {
69+ Self {
70+ data : HashMap :: new ( ) ,
71+ term : false ,
72+ }
73+ }
74+ fn insert ( & mut self , s : & [ u8 ] ) {
75+ if s. is_empty ( ) {
76+ return ;
77+ }
78+ if let Some ( node) = self . data . get_mut ( & s[ 0 ] ) {
79+ if s. len ( ) == 1 {
80+ node. as_mut ( ) . unwrap ( ) . term = true ;
81+ return ;
82+ }
83+ node. as_mut ( ) . unwrap ( ) . insert ( & s[ 1 ..] ) ;
84+ } else {
85+ let mut node = Self :: new ( ) ;
86+ if s. len ( ) == 1 {
87+ node. term = true ;
88+ } else {
89+ node. insert ( & s[ 1 ..] ) ;
90+ }
91+ self . data . insert ( s[ 0 ] , Some ( Box :: new ( node) ) ) ;
92+ }
93+ }
94+ fn find_prefix ( & self , s : & [ u8 ] , i : usize ) -> Vec < String > {
95+ if i == s. len ( ) {
96+ let mut prefix = Vec :: from ( s) ;
97+ return self . find_all ( & mut prefix) ;
98+ }
99+ if let Some ( node) = self . data . get ( & s[ i] ) {
100+ return node. as_ref ( ) . unwrap ( ) . find_prefix ( s, i+1 ) ;
101+ }
102+ return vec ! [ ]
103+ }
104+ fn find_all ( & self , prefix : & mut Vec < u8 > ) -> Vec < String > {
105+ let mut ans = vec ! [ ] ;
106+ if self . term {
107+ ans. push ( unsafe { std:: str:: from_utf8_unchecked ( & prefix) . to_string ( ) } ) ;
108+ }
109+ for ( & c, node) in self . data . iter ( ) {
110+ prefix. push ( c) ;
111+ let mut t = node. as_ref ( ) . unwrap ( ) . find_all ( prefix) ;
112+ if !t. is_empty ( ) {
113+ ans. append ( & mut t) ;
114+ }
115+ prefix. pop ( ) ;
116+ }
117+ ans
118+ }
119+ }
120+
121+ #[ cfg( test) ]
122+ mod tests {
123+ use super :: AutocompleteSystem ;
124+
125+ #[ test]
126+ fn test_input ( ) {
127+ let test_cases = vec ! [
128+ (
129+ vec![ "i love you" , "island" , "ironman" , "i love leetcode" ] ,
130+ vec![ 5 , 3 , 2 , 2 ] ,
131+ vec![
132+ ( 'i' , vec![ "i love you" , "island" , "i love leetcode" ] ) ,
133+ ( ' ' , vec![ "i love you" , "i love leetcode" ] ) ,
134+ ( 'a' , vec![ ] ) ,
135+ ( '#' , vec![ ] ) ,
136+ ( 'i' , vec![ "i love you" , "island" , "i love leetcode" ] ) ,
137+ ( ' ' , vec![ "i love you" , "i love leetcode" , "i a" ] ) ,
138+ ( 'a' , vec![ "i a" ] ) ,
139+ ( '#' , vec![ ] ) ,
140+ ( 'i' , vec![ "i love you" , "island" , "i a" ] ) ,
141+ ( ' ' , vec![ "i love you" , "i a" , "i love leetcode" ] ) ,
142+ ( 'a' , vec![ "i a" ] ) ,
143+ ( '#' , vec![ ] ) ,
144+ ] ,
145+ ) ,
146+ (
147+ vec![ "i love you" , "island" , "ironman" , "i love leetcode" ] ,
148+ vec![ 5 , 3 , 2 , 2 ] ,
149+ vec![
150+ ( 'i' , vec![ "i love you" , "island" , "i love leetcode" ] ) ,
151+ ( ' ' , vec![ "i love you" , "i love leetcode" ] ) ,
152+ ( 'l' , vec![ "i love you" , "i love leetcode" ] ) ,
153+ ( 'o' , vec![ "i love you" , "i love leetcode" ] ) ,
154+ ( 'v' , vec![ "i love you" , "i love leetcode" ] ) ,
155+ ( 'e' , vec![ "i love you" , "i love leetcode" ] ) ,
156+ ( ' ' , vec![ "i love you" , "i love leetcode" ] ) ,
157+ ( 'l' , vec![ "i love leetcode" ] ) ,
158+ ( 'c' , vec![ ] ) ,
159+ ( '#' , vec![ ] ) ,
160+ ( 'i' , vec![ "i love you" , "island" , "i love leetcode" ] ) ,
161+ ( ' ' , vec![ "i love you" , "i love leetcode" , "i love lc" ] ) ,
162+ ( 'l' , vec![ "i love you" , "i love leetcode" , "i love lc" ] ) ,
163+ ( 'o' , vec![ "i love you" , "i love leetcode" , "i love lc" ] ) ,
164+ ( 'v' , vec![ "i love you" , "i love leetcode" , "i love lc" ] ) ,
165+ ( 'e' , vec![ "i love you" , "i love leetcode" , "i love lc" ] ) ,
166+ ( ' ' , vec![ "i love you" , "i love leetcode" , "i love lc" ] ) ,
167+ ( 'y' , vec![ "i love you" ] ) ,
168+ ( '#' , vec![ ] ) ,
169+ ] ,
170+ ) ,
171+ ] ;
172+ for ( s, t, tc) in test_cases {
173+ let mut obj = AutocompleteSystem :: new ( s. iter ( ) . map ( |v| v. to_string ( ) ) . collect ( ) , t) ;
174+ for ( c, expect) in tc {
175+ assert_eq ! ( obj. input( c) , expect, "c:{}" , c) ;
176+ }
177+ }
178+ }
179+ }
0 commit comments