結果

問題 No.1324 Approximate the Matrix
ユーザー lam6er
提出日時 2025-03-20 20:51:38
言語 PyPy3
(7.3.15)
結果
WA  
実行時間 -
コード長 1,318 bytes
コンパイル時間 180 ms
コンパイル使用メモリ 82,424 KB
実行使用メモリ 86,096 KB
最終ジャッジ日時 2025-03-20 20:51:48
合計ジャッジ時間 6,027 ms
ジャッジサーバーID
（参考情報） 		judge5 / judge2
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 3
other AC * 9 WA * 33
権限があれば一括ダウンロードができます

ソースコード

diff # 

import heapq

def main():
    import sys
    input = sys.stdin.read().split()
    idx = 0
    N, K = int(input[idx]), 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]))
        idx += N
        P.append(row)
    
    sum_p_sq = sum(p * p for row in P for p in row)
    
    remaining_row = A.copy()
    remaining_col = B.copy()
    
    heap = []
    for i in range(N):
        for j in range(N):
            initial_cost = 1 - 2 * P[i][j]
            heapq.heappush(heap, (initial_cost, i, j))
    
    total_cost = 0
    x = [[0] * N for _ in range(N)]
    
    for _ in range(K):
        while True:
            if not heap:
                break  # should not happen since G is non-empty
            cost, i, j = heapq.heappop(heap)
            if remaining_row[i] > 0 and remaining_col[j] > 0:
                x[i][j] += 1
                remaining_row[i] -= 1
                remaining_col[j] -= 1
                total_cost += cost
                new_cost = 2 * x[i][j] + 1 - 2 * P[i][j]
                heapq.heappush(heap, (new_cost, i, j))
                break
    
    print(total_cost + sum_p_sq)

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