結果

問題 No.1324 Approximate the Matrix
ユーザー gew1fw
提出日時 2025-06-12 20:31:33
言語 PyPy3
(7.3.15)
結果
WA  
実行時間 -
コード長 2,735 bytes
コンパイル時間 191 ms
コンパイル使用メモリ 82,048 KB
実行使用メモリ 102,656 KB
最終ジャッジ日時 2025-06-12 20:32:22
合計ジャッジ時間 16,457 ms
ジャッジサーバーID
（参考情報） 		judge1 / 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]))
        P.append(row)
        idx +=N
    
    current_row = [0]*N
    current_col = [0]*N
    Q = [[0]*N for _ in range(N)]
    
    heap = []
    for i in range(N):
        for j in range(N):
            current_q = Q[i][j]
            marginal_cost = (current_q + 1 - P[i][j])**2 - (current_q - P[i][j])**2
            heapq.heappush(heap, (marginal_cost, i, j))
    
    remaining_A = A.copy()
    remaining_B = B.copy()
    
    while True:
        # Check if all are satisfied
        done = True
        for i in range(N):
            if remaining_A[i] >0:
                done = False
                break
        if done:
            for j in range(N):
                if remaining_B[j] >0:
                    done = False
                    break
        if done:
            break
        
        # Find the best possible cell to increment
        found = False
        temp_heap = []
        while heap and not found:
            mc, i, j = heapq.heappop(heap)
            if remaining_A[i] >0 and remaining_B[j] >0:
                # Check if current Q[i][j] is still the same (not been updated by others)
                # This part is a simplification and may not handle all cases correctly
                current_q = Q[i][j]
                expected_mc = (current_q +1 - P[i][j])**2 - (current_q - P[i][j])**2
                if mc == expected_mc:
                    # Proceed to increment
                    Q[i][j] +=1
                    remaining_A[i] -=1
                    remaining_B[j] -=1
                    # Push new marginal cost
                    new_mc = (Q[i][j] +1 - P[i][j])**2 - (Q[i][j] - P[i][j])**2
                    heapq.heappush(heap, (new_mc, i, j))
                    found = True
                else:
                    # This entry is outdated, discard
                    pass
                # Push back other entries to heap
                temp_heap.append((mc, i, j))
            else:
                temp_heap.append((mc, i, j))
        # Restore the heap
        for item in temp_heap:
            heapq.heappush(heap, item)
        if not found:
            # No valid cell found, which should not happen per problem statement
            break
    
    # Compute the result
    res =0
    for i in range(N):
        for j in range(N):
            res += (Q[i][j] - P[i][j])**2
    print(res)

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