【C++で実装】Swappable [ AtCoder Beginner Contest 206 ] 【初心者による初心者のための解説】

一昨日はじめてAtcoderに参加しました。参加したのは[ AtCoder Beginner Contest 206 ]です。結果はA、Bのみ正解できました。

C問題は本番で解けませんでしたが、大会終了後頑張って理解しました。今回はその内容を共有していきたいと思います。

なお、実装は公式解説の１を参考にしました。この解説は公式解説の１を読んだ前提で進めていきます。

公式解説：https://atcoder.jp/contests/abc206/editorial/2096

全体のコード

解説するコードは以下の通り。

#include <iostream>
#include <vector>
#include <algorithm>
#include <map>

using namespace std;

int main()
{
    long long N;
    cin >> N;

    vector<long long> A(N);

    for (int i = 0; i < N; ++i)
        cin >> A[i];

    long long count = ((N * (N - 1)) / 2);

    sort(A.begin(), A.end());

    map<int, long long> p_q;

    for (int i = 0; i < N - 1; ++i)
    {
        if (A[i] != A[i + 1])
            continue;
        else if (!p_q[A[i]]) // if empty, then init by 1.
            p_q[A[i]] = 1;
        p_q[A[i]]++;
    }

    for (pair<int, long long> element : p_q)
    {
        long long q = element.second;
        count -= ((q * (q - 1)) / 2);
    }

    cout << count << endl;
}

解説

以下の４つのステップで実装しました。題名が抽象的でスイマセン。

【Step１】組み合わせをすべて数える

上の図のようなイメージで考えるとわかりやすいと思います。この場合は７をiとして残りをjとしています。組み合わせの数（矢印の数）は９個ですね。

次に８をiとしてそれ以降の数をjとして考えると８矢印の数は８個になります。

これを繰り返していくと組み合わせの総数は９＋８＋７＋….＋１となります。これは交差１，初項９の等差数列です。なのでその和は項数をNとすると(N * (N – 1)) / 2で与えられます。

long long count = ((N * (N - 1)) / 2);

【Step２】配列をソートする

配列を小さいものから大きいものの順番にソートします。これはalgorithmをincludeしてsortを使えば実装できます。

sort(A.begin(), A.end());

【Step３】連続する数と、連続した回数を記録する

Aをソートすると以下のような配列が与えられます。

１が２回、８が３回、９が２回連続しています。

以下のコードによって、このような「連続する数と、その回数のペア」を求めることができます。

    map<int, long long> p_q;

    for (int i = 0; i < N - 1; ++i)
    {
        if (A[i] != A[i + 1])
            continue;
        else if (!p_q[A[i]]) 
            p_q[A[i]] = 1;
        p_q[A[i]]++;
    }

if文：i番目の数と、i+1番目の数が同じかどうか調べて、違かったらそれ以降の操作をする必要ないのでcontinueします。

else if文：この部分でははじめてp_qに値を登録するときに、1で初期化する処理をしています。

最後のとこ：i番目とi+1番目が等しく、かつすでに値がp_qに登録されている場合は単純に１を足します。

【Step４】Step３のデータをもとに答えから引く

まず、もう一度Step１の操作からどんな値を引きたいのか考えてみます。

例えば最初の８をiとして残りをjにして考えると、数えたいのは青色の矢印、数えたくないのは赤色の矢印になります。

２つ目の８をiとして残りをjとして数えると以下のようになります。

赤色の矢印は合計で2+1=3となります。

イメージが湧いたところで、赤色の矢印を数える方法を一般化してみましょう。

例えば７７７７という文字列の場合を考えます。最初の７を基準にして３本、２番めの７を基準にして２本、３番目の７を基準にして１本なので合計3+2+1です。

５５５５５の場合は同様に4+3+2+1本の赤色の矢印があることになります。

結局赤色の矢印の合計は初項１，交差１、項数ｑ-1の等差数列の和、すなわちq(q-1)/2になります。このｑはStep３で求めた値です。１の場合はｑが２、８は３、９は２です。

最後にStep１で求めた総数からそれぞれの赤い矢印を引いていきます。

count – (１の赤い矢印の合計) – (８の赤い矢印の合計) – (９の赤い矢印の合計)

↑のような感じで総数から引いていきます。

for (pair<int, long long> element : p_q)
    {
        long long q = element.second;
        count -= ((q * (q - 1)) / 2);
    }

おわりに（感想）

解説記事を書いてみたら想像以上に大変だった。