AtCoder Beginner Contest 194 A〜C問題

前書き

AtCoder Beginner Contest 194 A〜C問題を解きました。

A問題

難易度:灰
時間:7分

A,B = map(int,input().split())
if (A+B)>=15 and B>=8: print(1)
elif (A+B)>=10 and B>=3: print(2)
elif (A+B)>=3: print(3)
else: print(4)

B問題

難易度:灰
時間:RE→AC(40分)

atcoder.jp

時間かかりすぎた・・・RE出してしまった・・・。
うーん...あまり回答とは綺麗でないですが、一応ACとなったコードです。
作業AとBの中で最小時間で仕事できる従業員を探し、AとBが両方とも終了する時間を表示します。
一人の従業員はAとBを同時にこなせません。異なる従業員にそれぞれAとBを割り振ると、同時に処理できます。

こんな感じでパターン分けしました。

  • AB両方をこなす従業員iに任せるのが最短
  • Aと従業員i、Bを従業員jに任せるのが最短(それぞれAはi、Bはjが最短)
  • Aと従業員i、Bを従業員jに任せるのが最短(それぞれAB共にi or jが最短だが、両方こなすと最短でない場合)

追記:全探索で解けば一発だったみたいです。(全探索に弱い)
ただ私のやり方の方が計算量は少ないはず。コードは煩雑だが。
atcoder.jp

import numpy as np
N = int(input())
A,B = [],[]
_min_sum = 200000
for i in range(N):
  a,b = map(int,input().split())
  _min_sum = min(_min_sum,a+b)
  A.append([a,i])
  B.append([b,i])

A,B = sorted(A), sorted(B)
_minA,_minB = [], []
for j in range(N):
  if A[j][1]==B[j][1]:
    if A[j+1][1]<=B[j+1][1]: 
      _minA,_minB=A[j+1],B[j]
      break
    elif A[j+1][1]>B[j+1][1]:
      _minA,_minB=A[j],B[j+1]
      break
  else:
    _minA,_minB=A[j],B[j]
    break

print(min(max(_minA[0], _minB[0]),_min_sum))

C問題

時間:25分

atcoder.jp

普通に全探索で解くとTLE。
以下のように簡単に式を展開すると1重ループで解けます。

(i - j)**2 = i ** 2 - i* j + j**2

from itertools import accumulate
N = int(input())
A = list(map(int,input().split()))

b = list(accumulate(A))
result = 0
for j in range(N):
  result+=(N-1)*(A[j]**2)

for i in range(N-1):
  result-=2*A[i+1]*b[i]
print(result)