An Art Problem

    1#include <bits/stdc++.h>
    2using namespace std;
    3struct Node{
    4    int d, col, r, c;// d is distance, col = colorm, r= row, c = col
    5    // priority_queue is a max heap by default so we need to change it
    6    bool operator<(const Node &o) const {
    7        if (d != o.d) return d > o.d;// smaller distance has higher priority
    8        return col > o.col; // if tie, then smaller colour has higher priority
    9    }
    10};
    11vector<pair<int, int>> dir = {{0, 1}, {0, -1}, {-1, 0}, {1, 0}};
    12int main(){
    13    ios::sync_with_stdio(0); cin.tie(0);
    14    int n, m, k; cin >> n >> m >> k;
    15    vector<vector<int>> g(n, vector<int>(m));
    16    for (int i = 0; i < n; i++){
    17        for (int j = 0; j < m; j++){
    18            cin >> g[i][j];
    19        }
    20    }
    21    auto out = g;
    22    const int INF = 1e9;
    23    vector<vector<int>> dist(n, vector<int>(m, INF));
    24    vector<vector<int>> best(n, vector<int>(m, 0));
    25    priority_queue<Node> pq;
    26    for (int i = 0; i < n; i++){
    27        for (int j = 0; j < m; j++){
    28            if (g[i][j] != 0){
    29                dist[i][j] = 0;
    30                best[i][j] = g[i][j];
    31                pq.push({0, g[i][j], i, j});
    32            }
    33        }
    34    }
    35    // multi source bfs with distance up to k
    36    while (!pq.empty()){
    37        Node cur = pq.top(); pq.pop();
    38        if (cur.d > dist[cur.r][cur.c]) continue;
    39        if (cur.d == dist[cur.r][cur.c] && cur.col > best[cur.r][cur.c]) continue;
    40        for (auto [dr, dc] : dir){
    41            int nr = cur.r + dr, nc = cur.c + dc;
    42            if (nr >= 0 && nr < n && nc >= 0 && nc < m){
    43                int nd = cur.d + 1;
    44                if (nd > k) continue;
    45                if (nd < dist[nr][nc] || (nd == dist[nr][nc] && cur.col < best[nr][nc])){
    46                    dist[nr][nc] = nd;
    47                    best[nr][nc] = cur.col;
    48                    pq.push({nd, cur.col, nr, nc});
    49                }
    50            }
    51        }
    52    }
    53    for (int i = 0; i < n; i++){
    54        for (int j = 0; j < m; j++){
    55            if (g[i][j] == 0 && dist[i][j] <= k) out[i][j] = best[i][j];
    56        }
    57    }
    58    for (int i = 0; i < n; i++){
    59        for (int j = 0; j < m; j++){
    60            cout << out[i][j] << " ";
    61        }
    62        cout << '\n';
    63    }
    64    return 0;
    65}
    66