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

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

AtCoder ABC 152 D - Handstand 2 (400 点)

こういうのをスパッと解けるようになりたい

問題へのリンク

問題概要

 1 以上  N 以下の整数の組  (A, B) であって、

  •  A の最上位の値と  B の一の位の値が等しい
  •  B の最上位の値と  A の一の位の値が等しい

という条件を満たすものが何個あるかを求めよ。

制約

  •  1 \le N \le 2 \times 10^{5}

考えたこと

ペアリングの個数を数えよ、といわれたときの考察として、よくあるのが

  •  A を固定して、それに応じた  B が何個あるかを高速に数える

という方法だと思う。だが今回はこの方針で突き進もうとすると辛い気持ちになる。少し考えてみると


 A を固定したときの  B の個数は、 A の「最上位の値」と「一の位の値」のみに依存する


ということが見えてくる。 A = 12 であっても、 A = 17432 であっても、それに対応する  B の個数は等しい。よって、方針を少し変形して、以下のようにすれば OK!!!

  1. 「最上位の値  p」と「一の位の値  q」を固定したときの、それを満たす  1 以上  N 以下の整数の個数 num[  p ][  q ] を求める (この個数は  A でも  B でも等しい)

  2.  (p, q) の組としてありうる  9^{2} = 81 通りの組合せそれぞれについて、num[  p ][  q ] × num[  p ][  q ] の値を合算する

num[  p ][  q ] は、単純に  1 から  N までの値それぞれについて、「最上位の値」と「一の位の値」を求めてあげて、それに応じた num の値をインクリメントすれば OK。計算量は  O(N)

#include <iostream>
#include <vector>
#include <string>
using namespace std;
#define COUT(x) cout << #x << " = " << (x) << " (L" << __LINE__ << ")" << endl

int main() {
    int N;
    cin >> N;
    long long res = 0;
    vector<vector<long long>> num(10, vector<long long>(10, 0));
    for (int n = 1; n <= N; ++n) {
        vector<int> d;
        int nn = n;
        while (nn) {
            d.push_back(nn % 10);
            nn /= 10;
        }
        int a = d[0];
        int b = d.back();
        if (!a || !b) continue;
        ++num[a][b];
    }
    for (int a = 1; a < 10; ++a) {
        for (int b = 1; b < 10; ++b) {
            res += num[a][b] * num[b][a];
        }
    }
    cout << res << endl;
}