結果

問題 No.1324 Approximate the Matrix
コンテスト
ユーザー gew1fw
提出日時 2025-06-12 20:36:56
言語 PyPy3
(7.3.15)
結果
WA  
実行時間 -
コード長 1,953 bytes
コンパイル時間 465 ms
コンパイル使用メモリ 82,580 KB
実行使用メモリ 85,500 KB
最終ジャッジ日時 2025-06-12 20:37:22
合計ジャッジ時間 4,894 ms
ジャッジサーバーID
（参考情報） 		judge5 / judge3
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 3
other AC * 9 WA * 33
権限があれば一括ダウンロードができます

ソースコード

diff # 

import heapq

def main():
    import sys
    input = sys.stdin.read().split()
    ptr = 0
    N = int(input[ptr])
    ptr += 1
    K = int(input[ptr])
    ptr += 1
    
    A = list(map(int, input[ptr:ptr+N]))
    ptr += N
    B = list(map(int, input[ptr:ptr+N]))
    ptr += N
    
    P = []
    for _ in range(N):
        row = list(map(int, input[ptr:ptr+N]))
        P.append(row)
        ptr += N
    
    Q = [[0]*N for _ in range(N)]
    remaining_row = A.copy()
    remaining_col = B.copy()
    
    heap = []
    
    # Initialize the heap with all possible edges
    for i in range(N):
        for j in range(N):
            if remaining_row[i] > 0 and remaining_col[j] > 0:
                current_k = 0
                cost = 1 - 2 * P[i][j]  # For k=0, next cost is 1-2P
                heapq.heappush(heap, (cost, i, j, current_k))
    
    for _ in range(K):
        while True:
            if not heap:
                break  # This should not happen as per problem constraints
            cost, i, j, current_k = heapq.heappop(heap)
            # Check if this entry is still valid (current_k matches Q[i][j])
            if Q[i][j] != current_k:
                continue
            # Check if we can add a unit to this edge
            if remaining_row[i] <= 0 or remaining_col[j] <= 0:
                continue
            # Add the unit
            Q[i][j] += 1
            remaining_row[i] -= 1
            remaining_col[j] -= 1
            # Compute the next cost and push if possible
            new_k = Q[i][j]
            new_cost = 2 * new_k + 1 - 2 * P[i][j]
            if remaining_row[i] > 0 and remaining_col[j] > 0:
                heapq.heappush(heap, (new_cost, i, j, new_k))
            break
    
    # Calculate the result
    result = 0
    for i in range(N):
        for j in range(N):
            diff = Q[i][j] - P[i][j]
            result += diff * diff
    print(result)

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