結果
問題 |
No.1318 ABCD quadruplets
|
ユーザー |
![]() |
提出日時 | 2025-05-14 13:20:23 |
言語 | PyPy3 (7.3.15) |
結果 |
TLE
|
実行時間 | - |
コード長 | 4,126 bytes |
コンパイル時間 | 261 ms |
コンパイル使用メモリ | 82,788 KB |
実行使用メモリ | 101,560 KB |
最終ジャッジ日時 | 2025-05-14 13:21:21 |
合計ジャッジ時間 | 8,234 ms |
ジャッジサーバーID (参考情報) |
judge5 / judge3 |
(要ログイン)
ファイルパターン | 結果 |
---|---|
sample | AC * 3 |
other | AC * 13 TLE * 1 -- * 16 |
ソースコード
import sys # Set input function for potentially faster reading input = sys.stdin.readline def solve(): # Read N and M from input N, M = map(int, input().split()) # Use a dictionary to store the counts of pairs (s, q) # where s = a + b and q = a^2 + b^2 for 0 <= a, b <= M. counts_dict = {} # Iterate through all possible pairs (a, b) for a in range(M + 1): for b in range(M + 1): # Calculate sum s and sum of squares q s = a + b q = a * a + b * b # Create a tuple key (s, q) key = (s, q) # Increment the count for this key in the dictionary # Use dict.get(key, 0) to handle the first time a key is seen counts_dict[key] = counts_dict.get(key, 0) + 1 # Convert the dictionary items into a list of ((s, q), count) tuples # Iterating over a list might be slightly faster than iterating over dict items pairs_with_counts = list(counts_dict.items()) # Initialize the answer array `ans` of size N+1 with zeros. # ans[n] will store the value f(n, M). ans = [0] * (N + 1) # Get the number of distinct (s, q) pairs K = len(pairs_with_counts) # Calculate the threshold value 2*N. We only care about quadruplets (a, b, c, d) # such that the expression value n satisfies n <= N. # The condition is E(a, b, c, d) = n, which is equivalent to # (a+b+c+d)^2 + (a^2+b^2+c^2+d^2) = 2n. # Let S = a+b+c+d = s1+s2 and Q = a^2+b^2+c^2+d^2 = q1+q2. # We need S^2 + Q = 2n. So S^2 + Q <= 2N. threshold = 2 * N # Precompute squares up to the maximum possible sum S = 4*M # This avoids repeated calculations of S*S inside the loops. max_S = 4 * M squares = [k*k for k in range(max_S + 1)] # Iterate through all distinct pairs (s1, q1) with index i for i in range(K): # Unpack the key (s1, q1) and its count count1 (s1, q1), count1 = pairs_with_counts[i] # Consider the combination of pair i with itself ((a,b), (c,d) where (a,b) and (c,d) yield the same (s,q)) # This corresponds to quadruplets where (a+b, a^2+b^2) = (s1, q1) and (c+d, c^2+d^2) = (s1, q1). S_ii = s1 + s1 Q_ii = q1 + q1 # Check if S_ii is within the precomputed squares bounds (it should be) # Calculate val = S^2 + Q val_ii = squares[S_ii] + Q_ii # Check if the resulting value 2n satisfies 2n <= 2N if val_ii <= threshold: # Calculate n = (S^2 + Q) / 2. Use bit shift for integer division by 2. n_ii = val_ii >> 1 # Add count1 * count1 to ans[n_ii]. This accounts for choosing a pair type i twice. ans[n_ii] += count1 * count1 # Consider combinations of pair i with pair j where j > i # This avoids double counting pairs and handles combinations of distinct pair types. for j in range(i + 1, K): # Unpack the key (s2, q2) and its count count2 for pair j (s2, q2), count2 = pairs_with_counts[j] # Calculate S = s1 + s2 and Q = q1 + q2 S_ij = s1 + s2 Q_ij = q1 + q2 # Calculate val = S^2 + Q using precomputed squares val_ij = squares[S_ij] + Q_ij # Check if the resulting value 2n satisfies 2n <= 2N if val_ij <= threshold: # Calculate n = (S^2 + Q) / 2 n_ij = val_ij >> 1 # Add 2 * count1 * count2 to ans[n_ij]. # The factor of 2 accounts for the symmetric case (j, i). # Use bit shift for multiplication by 2. ans[n_ij] += (count1 * count2) << 1 # Prepare the output strings output_lines = [str(x) for x in ans] # Write the results to standard output, separated by newlines. # Add a final newline as required by the problem statement format. sys.stdout.write("\n".join(output_lines) + "\n") # Execute the solve function solve()