はじめに

最近競技プログラミングをはじめました。
計算量の考え方やアルゴリズムについて学べ、かつゲーム性があって非常に面白いです。
今回は、以下の問題で色々と学びを得たので記録しておく。(普段から競技プログラミングをしていた人からすると当たり前の知識だと思う)

atcoder.jp

C - Green Bin

実行時間制限: 2 sec / メモリ制限: 1024 MB
配点 : 300 点

問題文

文字列 $a$ に含まれる文字を何らかの順序で並べることで得られる文字列を $a$ のアナグラムと呼びます。例えば、greenbin は beginner のアナグラムです。このように、同じ文字が複数回現れるときはその文字をちょうどその回数だけ使わなければなりません。 $N$ 個の文字列 $s_1,s_2,…,s_N$ が与えられます。それぞれの文字列は長さが $10$ で英小文字からなり、またこれらの文字列はすべて異なります。二つの整数 i, j (1 $\leq$ i ＜ j $\leq$ N) の組であって、 $s_i$ が $s_j$ のアナグラムであるようなものの個数を求めてください。

制約

2 $\leq$ N $\leq$ $10^{5}$
$s_i$ は長さ 10 の文字列である。
$s_i$ の各文字は英小文字である。
$s_1, s_2, ..., s_N$ はすべて異なる。

やったこと

とりあえず以下のようなコードを書いた。

from collections import Counter
n = int(input())
S = [input() for _ in range(n)]
ans = 0
 
for i in range(n):
  S[i] = Counter(S[i])
 
for i in range(n):
  for j in range(i + 1, n):
    if S[i] == S[j]:
      ans += 1
 
print(ans)

入力をcollections.Counterメソッドに一度代入することで、オブジェクト同士の比較を可能にしてアナグラムか判定することができる。
しかし、2 $\leq$ N $\leq$ $10^{5}$ であるため、上記のコードだと2 秒という時間制限を過ぎてしまうことが判明。
アナグラムの判定は問題なさそうだが、文字列の比較を行う必要があるようです。
最終的に以下のコードで線形時間で動作し、解けることがわかった。

n = int(input())
sorted_word = {}

for i in range(n):
  word = ''.join(sorted(input()))
  sorted_word.setdefault(word, 0)
  sorted_word[word] += 1

def pair_num(n):
  return (n * (n - 1)) // 2

print(sum(map(pair_num, sorted_word.values())))

まず入力された文字列をソートし、ハッシュテーブルに格納します。
こうすることでアナグラムの文字列の組は、ソート後に同じ文字列になるので出現回数をカウントすることができるようになります。ちなみに、setdefaultメソッドは対象のオブジェクトに存在しないKeyを指定した場合は対象のオブジェクトに引数で指定したKeyとValueが設定され、対象のオブジェクトに引数で指定したKeyが存在する場合は対象オブジェクトの値を返すらしい(このとき対象のオブジェクト自体は変更されない)。
最後にアナグラムの文字列の出現回数をペアの数に変換します。
あるアナグラムの文字列をKeyするValueに格納された値をnとすると、アナグラムの組は $(n - 1) + (n - 2) + ... + (n - n)$ になることがわかります。
あとは等差数列の一般項に当てはめればペアの数も計算できます。
$\frac{n}{2} (a + l) = \frac{n}{2} ((n-1)+ 0) = \frac{1}{2} n(n - 1)$
最後に各アナグラムの文字列ごとのペアを合計して終わりです。