Line data    Source code

       1              : // Copyright (c) 2024 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              : 
      17              : #include <gtest/gtest.h>
      18              : 
      19              : #include <queue>
      20              : #include <string>
      21              : #include <unordered_map>
      22              : #include <vector>
      23              : 
      24              : namespace {
      25              : class Solution {
      26              : public:
      27            2 :     std::vector<std::string> topKFrequent(const std::vector<std::string>& words, int k) {
      28            2 :         std::unordered_map<std::string, int> mmp;
      29           18 :         for (const auto& word : words) {
      30           16 :             mmp[word]++;
      31              :         }
      32            2 :         auto my_compare = [](const std::pair<std::string, int>& p1,
      33              :                              const std::pair<std::string, int>& p2) {
      34            9 :             if (p1.second == p2.second) {
      35            2 :                 return p1.first < p2.first;
      36              :             }
      37            7 :             return p1.second > p2.second;
      38              :         };
      39              : 
      40              :         std::priority_queue<std::pair<std::string, int>, std::vector<std::pair<std::string, int>>,
      41              :                             decltype(my_compare)>
      42            2 :             queue(my_compare);
      43              : 
      44           10 :         for (auto& pair : mmp) {
      45            8 :             if (queue.size() < k) {
      46            6 :                 queue.push(pair);
      47              :             } else {
      48            2 :                 auto top = queue.top();
      49            4 :                 if ((top.second == pair.second && top.first > pair.first) ||
      50            2 :                     (top.first < pair.first)) {
      51            2 :                     queue.pop();
      52            2 :                     queue.push(pair);
      53              :                 }
      54            2 :             }
      55              :         }
      56            2 :         std::vector<std::string> ret;
      57            8 :         while (!queue.empty()) {
      58            6 :             ret.insert(ret.begin(), queue.top().first);
      59            6 :             queue.pop();
      60              :         }
      61            4 :         return ret;
      62            2 :     }
      63              : };
      64              : } // namespace
      65              : 
      66            4 : TEST(Leetcode, top_k_frequent_words) {
      67            1 :     Solution s;
      68              :     {
      69            5 :         std::vector<std::string> exp = {"i", "love"};
      70            9 :         EXPECT_EQ(exp, s.topKFrequent({"i", "love", "leetcode", "i", "love", "coding"}, 2));
      71            1 :     }
      72              : 
      73              :     {
      74            7 :         std::vector<std::string> exp = {"the", "is", "sunny", "day"};
      75           13 :         EXPECT_EQ(exp,
      76              :                   s.topKFrequent(
      77            1 :                       {"the", "day", "is", "sunny", "the", "the", "the", "sunny", "is", "is"}, 4));
      78            1 :     }
      79            9 : }