【AtCoder】ABC 345 C - One Time Swap | 茶コーダーが解くAtCoder

atcoder.jp

実行時間制限: 2 sec / メモリ制限: 1024 MB / Difficulty: 469

問題概要

文字列 $S$ に対して次の操作をちょうど1回行った後の文字列としてあり得るものの個数を求めよ。

操作 : $S$ の任意の2文字を選び、その2文字の位置を入れ替える。

制約

$S$ は英小文字からなる長さ $2$ 以上 $10^6$ 以下の文字列

考察

選んだ2文字の種類が同じ場合と異なる場合とで、場合分けしてみる。S = aabbbcdを例として考えよう。

同じ種類の場合

例えば、 $S$ の3文字目のbと5文字目のbを選んだとき。この場合、位置を入れ替えてできる文字列 $S'$ はaabbbcdとなり、 $S$ に一致する。これは1文字目のaと2文字目のaを選んだときなども同様である。

つまり、同じ種類の文字を選んだ場合は、その文字の種類に関わらず $S$ そのもの1通りが作り出せるのである。

異なる種類の場合

例えば、 $S$ の1文字目のaと3文字目のbを選んだとき。この場合、 $S'$ はbaabbcdとなる。また、 $S$ の2文字目のaと5文字目のbを選んだときは、 $S'$ はabbbacdとなる。

つまり、異なる種類の文字を選んだ場合は、相異なる $S'$ が作り出せる。

したがって主に数えたいのは、「 $S$ の中から異なる種類の2文字を選ぶ選び方は何通りあるか？」ということである。これは $S$ の長さを $N$ とすると全部で $\frac{N(N-1)}{2}$ 通りあるので、単純に2重ループで数え上げるのでは間に合わない。このうまい数え方は kyopro_freinds さんのユーザ解説がとても参考になったので詳細はそちらで。

ここまでくればあとは実装するだけ。私は文字の種類をkey、その文字が $S$ の中に含まれている個数をvalueとするmapを用いた。あるKeyのvalueが1つでも1より大きければ (=同じ種類の文字が選べる場合)、異なる種類の2文字を選ぶ選び方に1を加える必要があることに注意。

コード

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

#define rep(i, start, end) for (ll i = (start); i < (ll)(end); i++)

// ======================================== //

int main()
{
    string S;
    cin >> S;

    ll l = S.size();

    map<char, ll> mp;
    rep(i, 0, S.size()) mp[S[i]]++;

    ll ans = l * l;
    bool flag = false;
    for (auto itr = mp.begin(); itr != mp.end(); itr++)
    {
        ans -= (itr->second) * (itr->second);
        if (itr->second > 1)
            flag = true;
    }

    if (flag)
        cout << ans / 2 + 1 << endl;
    else
        cout << ans / 2 << endl;
}