竸プロ精進記 エイシングプログラミングコンテスト2020 C XYZ Triplets python 3

問題:https://atcoder.jp/contests/aising2020/tasks/aising2020_c

こういう系統の問題は、考えられる入力パターンを全て試し条件を満たすものを探し出す「全探索」が有効です。さらにこの問題の場合は少しひねられており効率よく回答するための閃きが要求されます。私はリアルタイム参加中、どうしても良い方法が思い浮かばず悔しい思いをしました。

コード例

N=int(input())
a=[0]*10**4
for j in range(1,101):
    for k in range(1,101):
        for h in range(1,101):
            tmp=h**2+j**2+k**2+h*j+j*k+k*h
            if tmp<=10000:
                a[tmp-1]+=1
for i in range(N):
    print(a[i])

3つの整数、j,k,hの3重のforループで全探索します。問題の制約上最大値が10000で、式に各整数の２乗の項が入っているため、それぞれ100ぐらいまでしらべれば良いでしょう。

各数値１から順番に調べていくときに、計算結果が10000以下ならば、最初に用意しておいた空のリストの該当するインデックスに数値を１足してカウントしていきます。各f(n)ごとに計算するのではなく、1回最大値まで計算すれば十分である、というのがこの問題のミソです。

最後に、作成したリストのうち必要な箇所を出力すれば完了です。