結果

問題 No.1324 Approximate the Matrix
ユーザー lam6er
提出日時 2025-04-16 16:35:26
言語 PyPy3
(7.3.15)
結果
WA  
実行時間 -
コード長 1,639 bytes
コンパイル時間 483 ms
コンパイル使用メモリ 81,728 KB
実行使用メモリ 84,300 KB
最終ジャッジ日時 2025-04-16 16:38:34
合計ジャッジ時間 5,841 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 = int(input[idx])
    idx += 1
    k = int(input[idx])
    idx += 1
    
    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
    
    sum_P_squared = sum(p*p for row in P for p in row)
    
    remaining_row = A.copy()
    remaining_col = B.copy()
    
    x = [[0]*n for _ in range(n)]
    
    heap = []
    
    # Initialize the heap with possible edges
    for i in range(n):
        for j in range(n):
            if remaining_row[i] > 0 and remaining_col[j] > 0:
                mc = 1 - 2 * P[i][j]
                heapq.heappush(heap, (mc, i, j))
    
    sum_marginal = 0
    
    for _ in range(k):
        while True:
            if not heap:
                break  # This should not happen as per problem statement
            mc, i, j = heapq.heappop(heap)
            if remaining_row[i] <= 0 or remaining_col[j] <= 0:
                continue
            # Assign this unit
            x[i][j] += 1
            sum_marginal += mc
            remaining_row[i] -= 1
            remaining_col[j] -= 1
            # Push next marginal cost if possible
            if remaining_row[i] > 0 and remaining_col[j] > 0:
                next_mc = 2 * (x[i][j] + 1) - 1 - 2 * P[i][j]
                heapq.heappush(heap, (next_mc, i, j))
            break
    
    answer = sum_marginal + sum_P_squared
    print(answer)

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