ABC143 B - TAKOYAKI FESTIVAL 2019

備忘録

問題

atcoder.jp

回答

回答.1

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()
    D = il()
    ret = 0
    for i in itertools.combinations(D, 2):
        ret += i[0] * i[1]
    print(ret)


if __name__ == '__main__':
    main()

回答.2

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()
    D = il()
    cum = itertools.accumulate(D)
    print(sum(c * d for c, d in zip(cum, D[1:])))


if __name__ == '__main__':
    main()

考え方

回答.1では、itertools.combinationsを使用して、たこ焼きのおいしさ(D)の
組み合わせをすべて網羅し、積（体力の回復量）を計算した。
単純に全ての組み合わせを試す回答方法。

回答.2では、itertools.accumulateを使用して、累積和を先に求める。
その後、配列Dの各要素に対して累積和との積（体力の回復量）を求める。

回答.2の考え方
例えばおいしさがa, b, c, dとあったとき、
おいしさaの各組み合わせから得られる回復量は、
a * b + a * c + a * d = a(b + c + d)となる。
同様に、おいしさbの組み合わせから得られる回復量は、
b * c + b * d = b(c + d)となる。
つまり、各a, b, c, dの組み合わせのおいしさを求めるときに、
先にカッコの中身を計算しておくことで、簡単に回答を求めることができる。

ちなみに回答.1の実行時間が21 ms、回答.2の実行時間が23 msだった。