C11 莫队 • EH's Blog

题目情境#

题目链接：https://www.luogu.com.cn/problem/P2709 ↗

题目描述

小B 有一个长为 $n$ 的整数序列 $a$ ，值域为 $[1,k]$ 。他一共有 $m$ 个询问，每个询问给定一个区间 $[l,r]$ ，求： $\sum\limits_{i=1}^k c_i^2$

其中 $c_i$ 表示数字 $i$ 在 $[l,r]$ 中的出现次数。小B请你帮助他回答询问。

输入格式

第一行三个整数 $n,m,k$ 。

第二行 $n$ 个整数，表示小B 的序列。

接下来的 $m$ 行，每行两个整数 $l,r$ 。

输出格式

输出 $m$ 行，每行一个整数，对应一个询问的答案。

输入 #1

6 4 3
1 3 2 1 1 3
1 4
2 6
3 5
5 6

plaintext

输出 #1

plaintext

对于 $100\%$ 的数据， $1\le n,m,k \le 5\times 10^4$ 。

算法解析#

Python代码实现#

import sys
import math
from collections import defaultdict
sys.setrecursionlimit(10**7)
input = sys.stdin.readline
n, m, k = map(int, input().split())
a = [0] + list(map(int, input().split()))

B = int(math.sqrt(n))

q = []
for i in range(m):
    l, r = map(int, input().split())
    q.append((l, r, i))

def mo_key(x):
    block = x[0] // B
    return (block, x[1] if block % 2 == 0 else -x[1])

q.sort(key=mo_key)

c = defaultdict(int)
sum = 0

def add(x):
    global sum
    sum -= c[x] * c[x]
    c[x] += 1
    sum += c[x] * c[x]

def remove(x):
    global sum
    sum -= c[x] * c[x]
    c[x] -= 1
    sum += c[x] * c[x]

ans = [0] * m
l, r = 1, 0
for L, R, idx in q:
    while l > L:
        l -= 1
        add(a[l])
    while r < R:
        r += 1
        add(a[r])
    while l < L:
        remove(a[l])
        l += 1
    while r > R:
        remove(a[r])
        r -= 1
    ans[idx] = sum

print("\n".join(map(str, ans)))

python

C++代码实现#

// 普通莫队 O(n*sqrt(n))
#include <iostream>
#include <cstring>
#include <algorithm>
#include <cmath>
using namespace std;

const int N=50005;
int n,m,k,B,a[N];
int sum,c[N],ans[N];
struct Q{
  int l,r,id;
  //按l所在块的编号l/B和r排序
  bool operator<(Q &b){
    if(l/B!=b.l/B) return l<b.l;
    if((l/B)&1) return r<b.r;
    return r>b.r;
  }
}q[N];

void add(int x){ //扩展一个数
  sum-=c[x]*c[x];
  c[x]++;
  sum+=c[x]*c[x];
}
void del(int x){ //删除一个数
  sum-=c[x]*c[x];
  c[x]--;
  sum+=c[x]*c[x];
}
int main(){
  scanf("%d%d%d",&n,&m,&k);
  B=sqrt(n); //块的大小
  for(int i=1;i<=n;++i)scanf("%d",&a[i]);
  for(int i=1;i<=m;++i)
    scanf("%d%d",&q[i].l,&q[i].r),q[i].id=i;
  sort(q+1,q+1+m); //按l/B和r排序
  for(int i=1,l=1,r=0;i<=m;++i){
    while(l>q[i].l) add(a[--l]);//左扩展
    while(r<q[i].r) add(a[++r]);//右扩展
    while(l<q[i].l) del(a[l++]);//左删除
    while(r>q[i].r) del(a[r--]);//右删除
    ans[q[i].id]=sum;
  }
  for(int i=1;i<=m;++i)printf("%d\n",ans[i]);
}