ABC159 復習

AtCoder ABC アルゴリズム Python

復習です。

A問題

ハマりました。
2つ選んで偶数になるのは、偶数・偶数か、奇数・奇数の組み合わせしかないので、N、Mそれぞれについて組み合わせを調べればヨシ。
pythonでは、itertoolsを使用すると解きやすいです。

import itertools
#例えば、10個のものから2個を選ぶ場合
ans = len(list(itertools.combinations(range(10), 2)))

B問題

文字操作について問われている問題ですね。
pythonでは、スライスをうまく利用して部分文字列とその逆順を求めればよいです。

#スライスの例
s = 'abcdef'
s1 = s[0:3]
s2 = s[3:]
#逆順
s3 = s[::-1]

C問題

ハマった。 入力例が偶然（？）にも3の倍数で例示されていることもあり、入力値/3の3乗でいいのでは？と導き出したところまではよかった。
小数計算はDecimalを使おうね、というお話。

D問題

TLEに阻まれました。
WAした提出コードはこれです。
もうね、2重ループはダメですよね。というわけで、ACしたコードがこちら（下記にも載せます）。

from collections import Counter
N = int(input())
A = list(map(int, input().split()))

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

counts = Counter(A)
total = 0
for n, count in counts.items(): 
    #初めに、同じ数同士から2つ選んだ時の組み合わせの総和を求める
    combi = choose2(count)
    total += combi 

for i in A:
    #上で求めた組み合わせの総和から、消える数字について1つ減らしたときの組み合わせを減らす
    temp = total - choose2(counts[i])
    temp += choose2(counts[i] - 1) 
    print(temp)

自分のために解説。
まず、ある数字を消すとかを考えずに同じ数字から2個選んだ時の組み合わせの総和を出しておきます。 この部分ですね。

counts = Counter(A)
total = 0
for n, count in counts.items(): 
    #初めに、同じ数同士から2つ選んだ時の組み合わせの総和を求める
    combi = choose2(count)
    total += combi

Counterは、dictのサブクラスで、イテラブルを渡すと文字や数字がいくつずつ表れているかを簡単に求めることができます。

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

あとはN個から2個を取り出す組み合わせについて、上記のコードを使用して、各組合せの総和を求めます。

最後に、下記の部分で変数totalから、N個の組み合わせとN-1個の組み合わせの差を引いてあげればそれが答えになります。

for i in A:
    #上で求めた組み合わせの総和から、消える数字について1つ減らしたときの組み合わせを減らす
    temp = total - choose2(counts[i])
    temp += choose2(counts[i] - 1) 
    print(temp)

余談

最初、組み合わせを求めるのにitertools.combinations()を使用していたのですが、TLEになってしまいました。
少し調べたところ、itertools.combinations()は組み合わせの 数 を返すものではなくて、組み合わせのタプルを返すメソッドなんですね。そりゃ遅くもなりますねぇ・・。