Line data    Source code

       1              : // Copyright (c) 2022 The Authors. All rights reserved.
       2              : //
       3              : // Licensed under the Apache License, Version 2.0 (the "License");
       4              : // you may not use this file except in compliance with the License.
       5              : // You may obtain a copy of the License at
       6              : //
       7              : //      https://www.apache.org/licenses/LICENSE-2.0
       8              : //
       9              : // Unless required by applicable law or agreed to in writing, software
      10              : // distributed under the License is distributed on an "AS IS" BASIS,
      11              : // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
      12              : // See the License for the specific language governing permissions and
      13              : // limitations under the License.
      14              : 
      15              : // Authors: liubang (it.liubang@gmail.com)
      16              : // Created: 2022/01/17 14:50
      17              : 
      18              : #include <gtest/gtest.h>
      19              : 
      20              : #include <queue>
      21              : #include <string>
      22              : #include <unordered_map>
      23              : #include <vector>
      24              : 
      25              : namespace {
      26              : class Solution {
      27              : public:
      28            2 :     std::vector<std::string> topKFrequent(const std::vector<std::string>& words, int k) {
      29            2 :         std::unordered_map<std::string, int> mmp;
      30           18 :         for (const auto& word : words) {
      31           16 :             mmp[word]++;
      32              :         }
      33            2 :         auto my_compare = [](const std::pair<std::string, int>& p1,
      34              :                              const std::pair<std::string, int>& p2) {
      35            9 :             if (p1.second == p2.second) {
      36            2 :                 return p1.first < p2.first;
      37              :             }
      38            7 :             return p1.second > p2.second;
      39              :         };
      40              : 
      41              :         std::priority_queue<std::pair<std::string, int>, std::vector<std::pair<std::string, int>>,
      42              :                             decltype(my_compare)>
      43            2 :             queue(my_compare);
      44              : 
      45           10 :         for (auto& pair : mmp) {
      46            8 :             if (queue.size() < k) {
      47            6 :                 queue.push(pair);
      48              :             } else {
      49            2 :                 auto top = queue.top();
      50            4 :                 if ((top.second == pair.second && top.first > pair.first) ||
      51            2 :                     (top.first < pair.first)) {
      52            2 :                     queue.pop();
      53            2 :                     queue.push(pair);
      54              :                 }
      55            2 :             }
      56              :         }
      57            2 :         std::vector<std::string> ret;
      58            8 :         while (!queue.empty()) {
      59            6 :             ret.insert(ret.begin(), queue.top().first);
      60            6 :             queue.pop();
      61              :         }
      62            4 :         return ret;
      63            2 :     }
      64              : };
      65              : } // namespace
      66              : 
      67            4 : TEST(Leetcode, top_k_frequent_words) {
      68            1 :     Solution s;
      69              :     {
      70            5 :         std::vector<std::string> exp = {"i", "love"};
      71            9 :         EXPECT_EQ(exp, s.topKFrequent({"i", "love", "leetcode", "i", "love", "coding"}, 2));
      72            1 :     }
      73              : 
      74              :     {
      75            7 :         std::vector<std::string> exp = {"the", "is", "sunny", "day"};
      76           13 :         EXPECT_EQ(exp,
      77              :                   s.topKFrequent(
      78            1 :                       {"the", "day", "is", "sunny", "the", "the", "the", "sunny", "is", "is"}, 4));
      79            1 :     }
      80            9 : }