ABC 159 D - Banned K

AtCoder Python3

備忘録

問題

atcoder.jp

回答

import sys
import os
import collections

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

    N = int(sys.stdin.buffer.readline().rstrip())
    A = list(map(int, sys.stdin.buffer.readline().split()))

    obj = collections.Counter(A)

    p = 0
    for v in obj.values():
        if v > 1:
            p += v * (v - 1) // 2

    for a in A:
        print(p - (obj[a] - 1))


if __name__ == '__main__':
    main()

考え方

ボールがN個の時の選べる組み合わせ数から、
上から順に整数列AのAiを除いたN-1個の時、選べる組み合わせ数を引くことで回答を得られる。

回答では、整数リストから、各要素（整数）の数え上げしたコレクションを作成している。

A = [1, 1, 2, 1, 2] # input
obj = collections.Counter(A) # {1: 3, 2: 2}

要素の個数がわかればnC2(n個から2個を選択）で組み合わせの数がわかるため、
全ての要素の組み合わせの合計が求めることができる。

あとは整数列Aを上から順に、合計 - （要素数 - 1)で、
Aiを除いた時のN-1個の時、選べる組み合わせ数を求めることができる。