题解：P17134 [KOI 2026 #1] 邻居

· · 题解

提示：该解法涉及知识点【前缀和】，难度为橙。

题目解法

虽然 n \le 3 \times 10^3，但我们还是有 O(n) 的解法。

做前缀和。定义 a_{i,j} 为前 i 个学生里属于学校 j 的人数，由于 s_i \in \{1,2\}，所以复杂度是 O(n) 的。

计算时，i 同校的邻居是在 [i-k_1,i+k_1] 中所有与 s_i 相同的，异校的邻居是在 [i-k_2,i+k_2] 中所有与 s_i 不同的，因此可以作差得到答案，再减去学生自己。

参考代码

由于 1 \oplus 2 = 3，所以 s_i \oplus 3 就是另一个学校。

#include<bits/stdc++.h>
using namespace std;
const int N=3001;
int n,k1,k2,s[N],a[N][3];
int main(){
    cin>>n>>k1>>k2;
    for(int i=1;i<=n;i++){
        cin>>s[i];
        a[i][1]=a[i-1][1];
        a[i][2]=a[i-1][2];
        a[i][s[i]]++;
    }
    for(int i=1;i<=n;i++){
        int sum=0,l,r;
        l=max(i-k1-1,0),r=min(i+k1,n);
        sum+=a[r][s[i]]-a[l][s[i]];
        s[i]^=3;
        l=max(i-k2-1,0),r=min(i+k2,n);
        sum+=a[r][s[i]]-a[l][s[i]];
        cout<<sum-1<<' ';
    }
    return 0;
}