ABC089 C - March

AtCoder Python3 collections itertools

備忘録

問題

atcoder.jp

回答

import sys
import os
import math
import bisect
import collections
import itertools
import heapq

ii = lambda: int(sys.stdin.buffer.readline().rstrip())
il = lambda: list(map(int, sys.stdin.buffer.readline().split()))
fl = lambda: list(map(float, sys.stdin.buffer.readline().split()))
iln = lambda n: [int(sys.stdin.buffer.readline().rstrip()) for _ in range(n)]

iss = lambda: sys.stdin.buffer.readline().decode().rstrip()
sl = lambda: list(map(str, sys.stdin.buffer.readline().decode().split()))
isn = lambda n: [sys.stdin.buffer.readline().decode().rstrip() for _ in range(n)]

lcm = lambda x, y: x * y / math.gcd(x, y)

MOD = 10 ** 9 + 7
MAX = float('inf')


def main():
    if os.getenv("LOCAL"):
        sys.stdin = open("input.txt", "r")

    N = ii()
    name = {}
    for al in range(ord('A'), ord('Z') + 1):
        name[chr(al)] = 0
    for n in range(N):
        name[iss()[0]] += 1

    print(sum([name[a] * name[b] * name[c] for a, b, c in itertools.combinations(['M', 'A', 'R', 'C', 'H'], 3)]))


if __name__ == '__main__':
    main()

考え方

赤のボール、青のボール、黄色のボール、各ボールが複数個ある場合、
同じ色の重複がないように３つのボールを選ぶ組み合わせは、
赤のボールの数 * 青のボールの数 * 黄色のボールの数となる。

この問題でも、上記の例と同様の考え方が出来るため、
対象のアルファベット('M', 'A' , 'R', 'C', 'H')から始まっている文字列の数をカウントし、
カウントを行った数から全ての組み合わせについて、何通りあるか求める。
ちなみに、各アルファベット('M', 'A' , 'R', 'C', 'H')から
3つを選択する重複のない組み合わせは、全10通り。

('M', 'A', 'R'), ('M', 'A', 'C'), ('M', 'A', 'H'), ('M', 'R', 'C'), ('M', 'R', 'H'), ('M', 'C', 'H'), ('A', 'R', 'C'), ('A', 'R', 'H'), ('A', 'C', 'H'), ('R', 'C', 'H')

実装はアルファベットをkey、 該当のアルファベットから始まる文字列の数をvalueとした辞書型を用意し、
文字列の数をカウントする。
その後、全ての組み合わせについて求め、結果の合算を回答とした。

collections.Counter()を使用すると、少しだけ簡潔に書くことができる。
回答例：
Submission #12933549 - AtCoder Beginner Contest 089