結果

問題 No.1318 ABCD quadruplets
ユーザー lam6er
提出日時 2025-03-31 17:44:04
言語 PyPy3
(7.3.15)
結果
WA  
実行時間 -
コード長 1,662 bytes
コンパイル時間 250 ms
コンパイル使用メモリ 82,584 KB
実行使用メモリ 85,888 KB
最終ジャッジ日時 2025-03-31 17:45:06
合計ジャッジ時間 5,304 ms
ジャッジサーバーID
(参考情報)
judge2 / judge1
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
sample WA * 3
other WA * 11 TLE * 1 -- * 18
権限があれば一括ダウンロードができます

ソースコード

diff #

import math
from collections import defaultdict

def main():
    import sys
    input = sys.stdin.read().split()
    N = int(input[0])
    M = int(input[1])

    # Precompute s_counts: a dictionary where key is s1 (a + b), value is another dict of t1 (a^2 + b^2) counts.
    s_counts = defaultdict(lambda: defaultdict(int))
    max_s_pair = 2 * M  # Maximum possible sum of a pair (a + b) which is M + M = 2M
    for a in range(M + 1):
        for b in range(M + 1):
            s = a + b
            t = a * a + b * b
            s_counts[s][t] += 1

    res = [0] * (N + 1)

    s_max = math.isqrt(2 * N)

    for s in range(s_max + 1):
        s_squared = s * s
        if s_squared > 2 * N:
            continue

        min_s1 = max(0, s - max_s_pair)
        max_s1 = min(s, max_s_pair)

        for s1 in range(min_s1, max_s1 + 1):
            s2 = s - s1
            if s2 < 0 or s2 > max_s_pair:
                continue

            # Check if s1 and s2 are in s_counts and have entries
            if not s_counts[s1] or not s_counts[s2]:
                continue

            for t1, cnt1 in s_counts[s1].items():
                for t2, cnt2 in s_counts[s2].items():
                    t_total = t1 + t2
                    two_n = s_squared + t_total
                    if two_n > 2 * N:
                        continue
                    if two_n % 2 != 0:
                        continue
                    n_val = two_n // 2
                    if 0 <= n_val <= N:
                        res[n_val] += cnt1 * cnt2

    # Output the result
    for n in range(N + 1):
        print(res[n])
    print()

if __name__ == "__main__":
    main()
0