結果
問題 |
No.1324 Approximate the Matrix
|
ユーザー |
![]() |
提出日時 | 2025-04-16 16:39:03 |
言語 | PyPy3 (7.3.15) |
結果 |
WA
|
実行時間 | - |
コード長 | 2,123 bytes |
コンパイル時間 | 231 ms |
コンパイル使用メモリ | 81,604 KB |
実行使用メモリ | 84,360 KB |
最終ジャッジ日時 | 2025-04-16 16:40:39 |
合計ジャッジ時間 | 4,814 ms |
ジャッジサーバーID (参考情報) |
judge3 / judge5 |
(要ログイン)
ファイルパターン | 結果 |
---|---|
sample | AC * 3 |
other | AC * 9 WA * 33 |
ソースコード
import heapq def main(): import sys input = sys.stdin.read().split() idx = 0 N = int(input[idx]) K = int(input[idx+1]) idx +=2 A = list(map(int, input[idx:idx+N])) idx +=N B = list(map(int, input[idx:idx+N])) idx +=N P = [] for _ in range(N): row = list(map(int, input[idx:idx+N])) P.append(row) idx +=N # Check if P is already valid valid = True for i in range(N): if sum(P[i]) != A[i]: valid = False break if valid: col_sums = [sum(row[j] for row in P) for j in range(N)] for j in range(N): if col_sums[j] != B[j]: valid = False break if valid: print(0) return # Initialize Q as all zeros, and adjust A and B to be the required row and column sums current_A = A.copy() current_B = B.copy() Q = [[0]*N for _ in range(N)] heap = [] # Precompute all possible edges and their initial marginal cost for i in range(N): for j in range(N): if current_A[i] > 0 and current_B[j] > 0: marginal_cost = 1 - 2 * P[i][j] # 2*0 +1 - 2*P[i][j] heapq.heappush(heap, (marginal_cost, i, j)) for _ in range(K): while True: if not heap: # This should not happen as per problem statement print(0) return mc, i, j = heapq.heappop(heap) if current_A[i] > 0 and current_B[j] > 0: break # Add this unit to Q[i][j] Q[i][j] +=1 current_A[i] -=1 current_B[j] -=1 # Push back the new marginal cost if possible if current_A[i] > 0 and current_B[j] > 0: new_mc = 2 * Q[i][j] +1 - 2 * P[i][j] heapq.heappush(heap, (new_mc, i, j)) # Compute the Frobenius norm squared total = 0 for i in range(N): for j in range(N): diff = Q[i][j] - P[i][j] total += diff * diff print(total) if __name__ == "__main__": main()