けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

第 4 回 PAST F - 構文解析 (4Q)

AtCoder PAST過去問 NoviSteps4Q 連想配列(setやmap) ソート ソート:ペア値 K番目を求める 集計処理

集計処理した上で、再度大きい順にソートする問題

問題概要

N 個の文字列 S_{1}, S_{2}, \dots, S_{N} が与えられる。

この列に一回以上出現する単語を、その出現回数の多い順に並べたとき K 番目の単語を出力せよ（一意に決まらない場合は "ambiguous" と出力せよ）。

制約

  • 1 \le N \le 10^{5}
  • 1 \le |S_{i}| \le 10

考えたこと

まずは集計処理しよう！　すなわち、次のような連想配列 (C++ ならば map<string, int> 型を求めよう。

  • nums[str]： N 個の文字列の中に、文字列 str が何個あるか

これを求めたあとは、各文字列ごとに再び (個数, 文字列) のペアを配列に格納していこう。これを大きい順にソートして K 番目を求めれば良い。

ただし、タイがある場合は "AMBIGUOUS" と返す必要がある。ここでは、次のように判定した。

  • K 番目に大きい頻度と、 K-1 番目に大きい頻度が一致する ( K-1 \ge 0 の場合)
  • K 番目に大きい頻度と、 K+1 番目に大きい頻度が一致する ( K+1 \gt 種類数 の場合)

については、 "AMBIGUOUS" と返せばよい。

全体として、文字列の長さの最大値を L として、計算量は O(LN \log N) と評価できる。

コード

#include <bits/stdc++.h>
using namespace std;

int main() {
    int N, K;
    cin >> N >> K, K--;
    string S;
    map<string, int> nums;
    for (int i = 0; i < N; i++) {
        cin >> S;
        nums[S]++;
    }

    vector<pair<int, string>> v;
    for (auto [str, num] : nums) v.emplace_back(num, str);
    sort(v.begin(), v.end());
    reverse(v.begin(), v.end());

    if (K > 0 && v[K].first == v[K-1].first) 
        cout << "AMBIGUOUS" << endl;
    else if (K+1 < v.size() && v[K].first == v[K+1].first)
        cout << "AMBIGUOUS" << endl;
    else
        cout << v[K].second << endl;
}