LeetCode Weekly Contest 381 : 3017. Count the Number of Houses at a Certain Distance II

久々の難易度8。
https://leetcode.com/contest/weekly-contest-381/problems/count-the-number-of-houses-at-a-certain-distance-ii/

問題

整数N,X,Yが与えられる。
1～NのN頂点、N辺からなるグラフがある。
i番と(i+1)番の頂点の間に計(N-1)本の辺があるほか、X番とY番の頂点の間に辺がある。

2頂点間の最短距離を考える。
その値が1～Nとなるのはそれぞれ何通りか。

解法

頂点uを総当たりし、より大きな頂点番号vに到達するための最短距離を考える。
vが1変わると最短距離は1変わるので、imos法の要領で区間に対し1加算するようにする。

uがYよりXに近い場合、X-Y間の辺を使うと距離が短縮される移動先があるので、その場所を尺取り法の要領で更新していくとよい。

class Solution {
public:
    vector<long long> countOfPairs(int n, int x, int y) {
		vector<ll> V(2*n+3);
		int i;
		if(x>y) swap(x,y);
		x--,y--;
		int d=abs(x-y);
		
		//内側から外側
		int t=x;
		FOR(i,n) {
			if(d<=1||abs(i-y)<=abs(i-x)) {
				V[0]+=2;
				V[n-i]-=2;
				continue;
			}
			while(t<y&&(t-i)<=((y-t)+abs(i-x)+1)) t++;
			if(t<=i) {
				V[0]+=2;
				V[n-i]-=2;
				continue;
			}
			V[0]+=2;
			V[t-i]-=2;
			V[abs(i-x)+1]+=2;
			V[abs(i-x)+2]+=2;
			V[abs(i-x)+(y-t)+2]-=2;
			V[abs(i-x)+(n-y)+1]-=2;
		}
		for(i=1;i<V.size();i++) V[i]+=V[i-1];
        V.erase(V.begin());
        V.resize(n);
        return V;
    }
};

まとめ

コードは短いし単純な問題設定だけど、最初からうまく書こうとすると割と戸惑う問題。