pairとmapを使いこなす練習 ABC152 D問題

Atcorder152 D問題

手も足も出なかった問題...
pairとmapを使えば一気に分かりやすくなった！
https://atcoder.jp/contests/abc152/tasks/abc152_d

#include<bits/stdc++.h>
#define rep(i,n) for(int i=0; i<(n); i++)
using namespace std;
using ll = long long;
const ll INF = 1e9 + 7;
typedef pair<int,int> P;// pair型は２つの値の組を表せる

P f(int x){
    int a=x%10; //下一桁
    int b =0;   //上一桁
    while(x){
        b=x;
        x/=10;
    }
    return P(a,b);
}

int main(){
    int n; cin >> n;
    map<P,int> freq; //連想配列の時に用いられる
    for (int i=1; i<=n; i++){
        P p=f(i);
        freq[p]++;
    }

    ll ans=0;
    for (int i=1; i<=n; i++){
        P p=f(i);
        P q(p.second,p.first);
        ans+=freq[q]; //A.firstとB.second, A.secondとB.firstが一致しているのか確認
    }
    cout << ans << endl;
    return 0;
}