結果

問題 No.1324 Approximate the Matrix
コンテスト
ユーザー lam6er
提出日時 2025-04-16 00:55:25
言語 PyPy3
(7.3.15)
結果
WA  
実行時間 -
コード長 2,123 bytes
コンパイル時間 194 ms
コンパイル使用メモリ 81,772 KB
実行使用メモリ 84,900 KB
最終ジャッジ日時 2025-04-16 00:57:23
合計ジャッジ時間 4,703 ms
ジャッジサーバーID
（参考情報） 		judge2 / judge3
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
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])
    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()
0