Snake Move

Snake Move

dijkstra，网格图

#include<bits/stdc++.h>
using namespace std;
using ll = long long;
using ull = unsigned long long;
using pii = pair<int, int>;
using tpi = tuple<int, int, int>;
const ll MOD = 998244353;
const int INF = 1e9;
// U D L R
const int dx[4] = {-1, 1, 0, 0};
const int dy[4] = {0, 0, -1, 1};

void solve() {
    int n, m, k; cin >> n >> m >> k;
    vector<vector<int>> grid(n + 1, vector<int>(m + 1, 0)); // 0代表'.'  -1代表'#'    其他数字代表从尾到头的编号
    int headx = -1, heady = -1;
    int neckx = -1, necky = -1;
    for(int i = 0; i < k; i ++) {
        int x, y; cin >> x >> y;
        if(i == 0) headx = x, heady = y;
        else if(i == 1) neckx = x, necky = y;
        grid[x][y] = k - i;
    }
    for(int i = 1; i <= n; i ++) {
        string s; cin >> s;
        for(int j = 1; j <= m; j ++) {
            if(s[j - 1] == '#') grid[i][j] = -1;
        }
    }
    
    priority_queue<tpi, vector<tpi>, greater<tpi>> q; //w, x, y
    vector<vector<int>> dist(n + 1, vector<int>(m + 1, INF));
    q.push({0, headx, heady}); dist[headx][heady] = 0; 
    if(neckx != -1) {
        dist[neckx][necky] = k - 1;
        q.push({k - 1, neckx, necky});
    }

    while(!q.empty()) {
        auto [w, x, y] = q.top(); q.pop();
        if(dist[x][y] < w) continue;
        for(int i = 0; i < 4; i ++) {
            int tx = x + dx[i], ty = y + dy[i];
            if(tx <= 0 || ty <= 0 || tx > n || ty > m) continue;
            if(grid[tx][ty] == -1) continue;
            if(dist[tx][ty] <= w + 1) continue;

            if(grid[tx][ty] > 0) {
                int tw = (w >= grid[tx][ty] - 1 ? w + 1 : grid[tx][ty]); // w过小，尾巴没收完，可以通过操作S缩短
                dist[tx][ty] = tw;
                q.push({tw, tx, ty});
            }
            else if(dist[tx][ty] > w + 1) {
                dist[tx][ty] = w + 1;
                q.push({w + 1, tx, ty});
            }
        }
    }

    ull ans = 0;
    for(int i = 1; i <= n; i ++) {
        for(int j = 1; j <= m; j ++) {
            if(dist[i][j] != INF) ans += (ull)dist[i][j] * dist[i][j];
        }
    }
    cout << ans << "\n";
}

int main() {
    ios::sync_with_stdio(false); cin.tie(0); cout.tie(0);
    int _ = 1; //cin >> _;
    while(_--) solve();
    return 0;
}