ARC 102 / C - Triangular Relationship

C - Triangular Relationship

過去問潰そうとしてみたけどいきなり結構厳しい問題で詰まった

このぐらいはサッと解けたいのでじっくり考える

全列挙する

とりあえず全列挙書いちゃダメなのはパッと見でわかるけど、
まずは書いて法則見つけないと全くわからん
ので一度全列挙型のダメパターンを書く

n,k = map(int,input().split())
import itertools
flag = True
ans=0
for i in itertools.product(list(range(1,n+1)),repeat=3):
    for j in itertools.permutations(i,2):
        if sum(j) % k != 0:
            flag = False
    if flag == True:
        ans +=1
    flag = True
print(ans)

productで重複あり組み合わせ。
permutationsで所謂普通の組み合わせを列挙。
flagをスイッチしてansをプラスしていく、普通の書き方。
このコードで値を入れて書き出すと
とりあえず問題の答え自体は分かるので法則性を見つける。

n	k	answer
3	2	9
3	3	1
4	2	16
4	3	1
5	1	125
5	2	35
5	3	1
5	4	2
6	1	216

kの偶奇見て平方数みたいな数出せばいいんじゃね？
と思ったけど
3x3 =9
4x4 =16
5x5x5 = 125 ??
どういうこっちゃ
問題名から三角数みたいな数列をイメージしてたけど違うっぽい
この時点で実際のコンテストなら分からなくて詰む気がした。

なので、ググって解説だけ見ながら一度それをコードにしてみた。

解法①

n,k = map(int,input().split())
ans1 = 0
ans2 = 0
for x in range(1,n+1):
    if x % k ==0:
        ans1 +=1
if k % 2 == 0:
    for y in range(1,n+1):
        if y % k == k/2:
            ans2 +=1
print(ans1**3+ans2**3)

で、三乗の足し算ってのは分かったけど、
そもそもfor文とか一切いらないのでは？と思った。
①kが偶数か奇数か判定する
②奇数の場合はnをkで割った商を三乗したのが答え
③偶数の場合はnをkで割った商を三乗+nをkで割った値を四捨五入して三乗

n=3,k=2なら
1+(1.5の四捨五入で2の三乗)=9
なんかこれで答えでそうだし計算量全然少ないかも

解法②

from decimal import Decimal, ROUND_HALF_UP
n,k = map(int,input().split())
ans = 0
if k % 2 ==0:
    p = Decimal(n/k).quantize(Decimal('0'),rounding=ROUND_HALF_UP)
    ans = p**3 +(n//k)**3
else:
    ans = (n // k)**3
print(ans)

めっちゃシンプルな計算でAC
自分的には②の方がシックリ来るけど、思いつかなきゃ即死
①もかけた方が圧倒的にいい気がする。
おそらく漸化式みたいな形で解く解法
公式の解説見る限り四捨五入の所は微妙な予感、もうちょっといい表現方法ありそう
こういうの正確に素早く解くって相当難しい気がする
atcoderこういうのサクっとつぶしている人がワラワラ居そうで怖い場所です。