Kyatatsu
Published on

包除原理を使ってD次元累積和を考える

Authors
  • Name
    Twitter

atcoderで3次元累積和の問題が出て,一般化できないか考えました. 役には立たないと思います.

とりあえずatcoderの問題(3次元累積和)をおいておきます.

問題の定式化

まずは通常の累積和(1次元累積和)を考えます.

1次元累積和

正整数 NN, 1xN1 \leq x \leq N を満たす整数 xx に対して 整数 AxA_x が与えられます. QQ 個のクエリが与えられるので,それぞれに答えてください.

qq 個目(1qQ1 \leq q \leq Q)のクエリでは整数 ll, rr (1lrN1 \leq l \leq r \leq N) が与えられるので,以下の値を求めてください.

S=x=lrAxS = \sum_{x=l}^{r} A_x

制約

  • 1N2×1051 \leq N \leq 2 \times 10^5
  • 1Q2×1051 \leq Q \leq 2 \times 10^5
  • 1Ax1091 \leq A_x \leq 10^9
  • 1lrN1 \leq l \leq r \leq N

解法

各クエリに対して x=lrAx\sum_{x=l}^{r} A_x を求めることが思い浮かびますが,これではクエリごとに計算量 O(N)O(N)となり,全体で O(NQ)O(NQ) なので間に合いません.

ここで以下の数列を定義します.

Si=k=1iAkS_i = \sum_{k=1}^{i} A_k

仮にこの数列が求まっているとすれば,クエリ l,rl, rに対して

x=lrAx=x=1rAix=1l1Ax=SrSl1\sum_{x=l}^{r} A_x = \sum_{x=1}^{r} A_i - \sum_{x=1}^{l-1} A_x = S_r - S_{l-1}

なので,O(1)O(1)SS を求めることができます.

また,SiS_iは漸化式 Si=Si1+AiS_i = S_{i-1} + A_i によって O(N)O(N)で求めることができるので,これを前計算しておけば,全体で O(N+Q)O(N+Q) でこの問題を解くことができます.

2次元累積和

次に2次元累積和を考えます.

正整数 NN, 1x,yN1 \leq x, y \leq N を満たす整数 xxyy に対して 整数 Ax,yA_{x, y} が与えられます.

QQ 個のクエリが与えられるので,それぞれに答えてください.

qq 個目(1qQ1 \leq q \leq Q)のクエリでは整数 lx,rx,ly,ryl_x, r_x, l_y, r_y (1lxrxN1 \leq l_x \leq r_x \leq N, 1lyryN1 \leq l_y \leq r_y \leq N) が与えられるので,以下の値を求めてください.

S=x=lxrxy=lyryAx,yS = \sum_{x=l_x}^{r_x} \sum_{y=l_y}^{r_y} A_{x, y}

制約

  • 1N30001 \leq N \leq 3000
  • 1Q30001 \leq Q \leq 3000
  • 1Ax,y1091 \leq A_{x, y} \leq 10^9
  • 1lxrxN1 \leq l_x \leq r_x \leq N
  • 1lyryN1 \leq l_y \leq r_y \leq N

解法

この問題に対しても1次元累積和と同じように,以下の数列を定義します.

Si,j=x=1iy=1jAx,yS_{i, j} = \sum_{x=1}^{i} \sum_{y=1}^{j} A_{x, y}

この数列が求まっているとすれば,クエリ lx,rx,ly,ryl_x, r_x, l_y, r_yに対して

x=lxrxy=lyryAx,y=Srx,rySlx1,rySrx,ly1+Slx1,ly1\sum_{x=l_x}^{r_x} \sum_{y=l_y}^{r_y} A_{x, y} = S_{r_x, r_y} - S_{l_x-1, r_y} - S_{r_x, l_y-1} + S_{l_x-1, l_y-1}

を計算すれば良いです.

また,Si,jS_{i, j}は漸化式 Si,j=Si1,j+Si,j1Si1,j1+Ai,jS_{i, j} = S_{i-1, j} + S_{i, j-1} - S_{i-1, j-1} + A_{i, j} によって O(N2)O(N^2)で求めることができるので,これを前計算しておけば,全体で O(N2+Q)O(N^2+Q) でこの問題を解くことができます.

この辺が怪しい人は以下の記事を読んでみてください.

ここで, x=lxrxy=lyryAx,y\sum_{x=l_x}^{r_x} \sum_{y=l_y}^{r_y} A_{x, y} を計算する際に Slx1,ly1S_{l_x-1, l_y-1} を足しているのは,Slx1,ly1S_{l_x-1, l_y-1}Slx1,ryS_{l_x-1, r_y}Srx,ly1S_{r_x, l_y-1} で重複している部分を差し引くためです.

2次元累積和のイメージ

  1. まず範囲全体を足す.
  2. Slx,ry1S_{l_x, r_y-1} を引く.
  3. 同様に Srx1,lyS_{r_x-1, l_y} を引く. ここで,Slx1,ly1S_{l_x-1, l_y-1} を2回引いていることになる.
  4. 引きすぎてしまったSlx1,ly1S_{l_x-1, l_y-1} を足す.

集合だと思ってみる

Slx,ry1S_{l_x, r_y-1}の領域を S1S_1Srx1,lyS_{r_x-1, l_y}の領域を S2S_2Slx1,ly1S_{l_x-1, l_y-1}の領域を S3S_3 とすると,

S3=S1S2S_3 = S_1 \cap S_2

となっていることがわかります.

また,2次元の和を求める際の式を少し変形すると,包除原理が使えそうなことがわかります.

x=lxrxy=lyryAx,y=Srx,rySlx1,rySrx,ly1+Slx1,ly1=Srx,ry(Slx1,ry+Srx,ly1Slx1,ly1)=Srx,ry(S1+S2S3)=Srx,ryS1S2\begin{align} \sum_{x=l_x}^{r_x} \sum_{y=l_y}^{r_y} A_{x, y} &= S_{r_x, r_y} - S_{l_x-1, r_y} - S_{r_x, l_y-1} + S_{l_x-1, l_y-1} \notag \\ &= S_{r_x, r_y} - (S_{l_x-1, r_y} + S_{r_x, l_y-1} - S_{l_x-1, l_y-1}) \notag \\ &= S_{r_x, r_y} - (S_1 + S_2 - S_3) \notag \\ &= S_{r_x, r_y} - |S_1 \cup S_2| \notag \end{align}

数学的に正しい記述方法がわかりませんが,気持ちは伝わると思います.

3次元累積和

これは一般に DD次元累積和に拡張できます. 例えば3次元の和を求めるには,

x=lxrxy=lyryz=lzrzAx,y,z=Srx,ry,rzS1S2S3\begin{align} \sum_{x=l_x}^{r_x} \sum_{y=l_y}^{r_y} \sum_{z=l_z}^{r_z} A_{x, y, z} &= S_{r_x, r_y, r_z} - |S_1 \cup S_2 \cup S_3| \notag \end{align}

とすればよいです. ここで,S1,S2,S3S_1, S_2, S_3 はそれぞれ Slx1,ry,rz,Srx,ly1,rz,Srx,ry,lz1S_{l_x-1, r_y, r_z}, S_{r_x, l_y-1, r_z}, S_{r_x, r_y, l_z-1} の領域を表しています.

しかし,具体的な S1S2S3|S_1 \cup S_2 \cup S_3| の表示が得られておらず,このままでは計算することができません.

そこで3変数の包除原理を考えます.S1S2=TS_1 \cup S_2 = Tとおきます.

S1S2S3=TS3|S_1 \cup S_2 \cup S_3| = |T \cup S_3|

2変数の包除原理を用いて,

TS3=T+S3TS3\begin{align} |T \cup S_3| = |T| + |S_3| - |T \cap S_3| \end{align}

TS3|T \cap S_3|について, 分配法則より

TS3=(S1S2)S3=(S1S3)(S2S3)=S1S3+S2S3(S1S3)(S2S3)=S1S3+S2S3S1S2S3\begin{align} |T \cap S_3| &= |(S_1 \cup S_2) \cap S_3| \notag \\ &= |(S_1 \cap S_3) \cup (S_2 \cap S_3)| \notag \\ &= |S_1 \cap S_3| + |S_2 \cap S_3| - |(S_1 \cap S_3) \cap (S_2 \cap S_3)| \notag \\ &= |S_1 \cap S_3| + |S_2 \cap S_3| - |S_1 \cap S_2 \cap S_3| \end{align}

となります.

また,T|T|について,

T=S1S2=S1+S2S1S2\begin{align} |T| &= |S_1 \cup S_2| \notag \\ &= |S_1| + |S_2| - |S_1 \cap S_2| \end{align}

となり,(2),(3)(2), (3)(1)(1)に代入すると,

S1S2S3=S1+S2+S3S1S2S1S3S2S3+S1S2S3\begin{align} |S_1 \cup S_2 \cup S_3| &= |S_1| + |S_2| + |S_3| \notag \\ &- |S_1 \cap S_2| - |S_1 \cap S_3| - |S_2 \cap S_3| \notag \\ &+ |S_1 \cap S_2 \cap S_3| \end{align}

が得られます.

ここで,

S1=Slx1,ry,rzS2=Srx,ly1,rzS3=Srx,ry,lz1|S_1| = S_{l_x-1, r_y, r_z} \\ |S_2| = S_{r_x, l_y-1, r_z} \\ |S_3| = S_{r_x, r_y, l_z-1} \\

であり,

S1S2=Slx1,ly1,rzS1S3=Slx1,ry,lz1S2S3=Srx,ly1,lz1S1S2S3=Slx1,ly1,lz1|S_1 \cap S_2| = S_{l_x-1, l_y-1, r_z} \\ |S_1 \cap S_3| = S_{l_x-1, r_y, l_z-1} \\ |S_2 \cap S_3| = S_{r_x, l_y-1, l_z-1} \\ |S_1 \cap S_2 \cap S_3| = S_{l_x-1, l_y-1, l_z-1} \\

となります.

なぜならば,例えば S1S2|S_1 \cap S_2|について,S1S2S_1 \cap S_2Slx1,ry,rzS_{l_x-1, r_y, r_z}Srx,ly1,rzS_{r_x, l_y-1, r_z} の共通した領域を表しており,zz方向は rzr_zと等しく,x,yx, y方向はそれぞれ狭い方の領域に"縮む"ため,それぞれlx1,ly1l_x-1, l_y-1となります.

D次元累積和

それではいよいよ一般化します. まず問題を定式化します.

正整数 NN, 1x1,x2,,xDN1 \leq x_1, x_2, \ldots, x_D \leq N を満たす整数 x1,x2,,xDx_1, x_2, \ldots, x_D に対して 整数 Ax1,x2,,xDA_{x_1, x_2, \ldots, x_D} が与えられます.

QQ 個のクエリが与えられるので,それぞれに答えてください.

qq 個目(1qQ1 \leq q \leq Q)のクエリでは整数 l1,r1,l2,r2,,lD,rDl_1, r_1, l_2, r_2, \ldots, l_D, r_D (1liriN1 \leq l_i \leq r_i \leq N) が与えられるので,以下の値を求めてください.

S=x1=l1r1x2=l2r2xD=lDrDAx1,x2,,xDS = \sum_{x_1=l_1}^{r_1} \sum_{x_2=l_2}^{r_2} \ldots \sum_{x_D=l_D}^{r_D} A_{x_1, x_2, \ldots, x_D}

解法

例によって累積和 Si1,i2,,iDS_{i_1, i_2, \ldots, i_D}O(ND)O(N^D)で求めておきます.

クエリ l1,r1,l2,r2,,lD,rDl_1, r_1, l_2, r_2, \ldots, l_D, r_Dに対して,

x1=l1r1x2=l2r2xD=lDrDAx1,x2,,xD=Sr1,r2,,rDS1S2SD\sum_{x_1=l_1}^{r_1} \sum_{x_2=l_2}^{r_2} \ldots \sum_{x_D=l_D}^{r_D} A_{x_1, x_2, \ldots, x_D} = S_{r_1, r_2, \ldots, r_D} - |S_1 \cup S_2 \cup \ldots \cup S_D|

を計算すればよいです.

もちろんS1,S2,,SDS_1, S_2, \ldots, S_D はそれぞれ Sl11,r2,,rD,Sr1,l21,,rD,,Sr1,r2,,lD1S_{l_1-1, r_2, \ldots, r_D}, S_{r_1, l_2-1, \ldots, r_D}, \ldots, S_{r_1, r_2, \ldots, l_D-1} の領域を表しています.

S1S2SD|S_1 \cup S_2 \cup \ldots \cup S_D|については,DD変数の包除原理

S1S2SD=i=1DSi1i<jDSiSj+1i<j<kDSiSjSk+(1)D1S1S2SD\begin{align} |S_1 \cup S_2 \cup \ldots \cup S_D| &= \sum_{i=1}^{D} |S_i| \notag \\ &- \sum_{1 \leq i < j \leq D} |S_i \cap S_j| \notag \\ &+ \sum_{1 \leq i < j < k \leq D} |S_i \cap S_j \cap S_k| \notag \\ &- \ldots \notag \\ &+ (-1)^{D-1} |S_1 \cap S_2 \cap \ldots \cap S_D| \end{align}

によって求めることができます. 証明については省略しますが,数学的帰納法を用いることで示すことができます.

急ですがここで,SSの添字をビット表現にすることを考えます.具体的には,整数i(1iD)i (1 \leq i \leq D)に対して,SSii次元目の添字はli,ril_i, r_iのどちらかですが,これがrir_iのときiiビット目を立てることにします.

例えば,S10110S_{10110}と表記したら,Sr1,l21,r3,r4,l51S_{r_1, l_2-1, r_3, r_4, l_5-1}を表すこととします.

このように表現することで,ビットごとのAND演算を用いて集合の共通部分を簡単に求められるようになります.

例えば,S1S3S4|S1 \cap S3 \cap S4|を求めてみましょう.

S1=Sl11,r2,r3,r4,r5=S01111S3=Sr1,r2,l31,r4,r5=S11011S4=Sr1,r2,r3,l41,r5=S11101S_1 = S_{l_1-1, r_2, r_3, r_4, r_5} = S_{01111} \\ S_3 = S_{r_1, r_2, l_3-1, r_4, r_5} = S_{11011} \\ S_4 = S_{r_1, r_2, r_3, l_4-1, r_5} = S_{11101}

なので,

S1S3S4=S01111&S11011&S11101=S01001=Sl11,r2,l31,l41,r5\begin{align} S_1 \cap S_3 \cap S_4 &= S_{01111} \& S_{11011} \& S_{11101} \notag \\ &= S_{01001} \notag \\ &= S_{l_1-1, r_2, l_3-1, l_4-1, r_5} \notag \end{align}

とわかります.

ちなみに,包除原理を適用した際の符号に関してもビットのpop countの偶奇を見ればわかります. これは,(5)において奇数個の共通部分は++,偶数個の共通部分は-であることから従います.

このように定式化したことで,例えば3次元累積和は以下のようなコードで実装できます.

#include<bits/stdc++.h>
using namespace std;
using ll = long long;

int N, Q;
int A[110][110][110], S[110][110][110];
int L[3], R[3];
int solve() {
    int sum = 0;
    for(int b=0; b<1<<3; b++) {
      int ind[3];
      for (int i=0; i<3; i++) {
          ind[i] = ((b & (1 << i)) ? L[i] - 1 : R[i]);
      }
      sum += S[ind[0]][ind[1]][ind[2]] * (__builtin_parity(b) ? -1 : 1);
    }
    return sum;
}

int main() {
    cin.tie(nullptr);
    ios::sync_with_stdio(false);
    cin>>N;
    for(int i=0; i<N; i++)for(int j=0; j<N; j++)for(int k=0; k<N; k++) cin>>A[i][j][k];
    for(int i=1; i<=N; i++)for(int j=1; j<=N; j++)for(int k=1; k<=N; k++) {
      L[0] = i; L[1] = j; L[2] = k;
      R[0] = i; R[1] = j; R[2] = k;
      S[i][j][k] = A[i-1][j-1][k-1] - solve();
    }
    cin>>Q;
    while(Q--) {
        for(int i=0; i<3; i++) cin>>L[i]>>R[i];
        cout<<solve()<<endl;
    }
}

これでDD次元の累積和を解くことができます! 計算量はO(ND+2DQ)O(N^D + 2^D Q)です.

おまけ

D=4D=4についてやってみましょう. 未来のABCで出題されるかもしれません.

x1=l1r1x2=l2r2x3=l3r3x4=l4r4Ax1,x2,x3,x4=Sr1,r2,r3,r4S1S2S3S4=Sr1,r2,r3,r4S1S2S3S4+S1S2+S1S3+S1S4+S2S3+S2S4+S3S4S1S2S3S1S2S4S1S3S4S2S3S4+S1S2S3S4=Sr1,r2,r3,r4Sl11,r2,r3,r4Sr1,l21,r3,r4Sr1,r2,l31,r4Sr1,r2,r3,l41+Sl11,l21,r3,r4+Sl11,r2,l31,r4+Sl11,r2,r3,l41+Sr1,l21,l31,r4+Sr1,l21,r3,l41+Sr1,r2,l31,l41Sl11,l21,l31,r4Sl11,l21,r3,l41Sl11,r2,l31,l41Sr1,l21,l31,l41+Sl11,l21,l31,l41\begin{align} & \sum_{x_1=l_1}^{r_1} \sum_{x_2=l_2}^{r_2} \sum_{x_3=l_3}^{r_3} \sum_{x_4=l_4}^{r_4} A_{x_1, x_2, x_3, x_4} \notag \\ &= S_{r_1, r_2, r_3, r_4} - |S_1 \cup S_2 \cup S_3 \cup S_4| \notag \\ &= S_{r_1, r_2, r_3, r_4} - |S_1| - |S_2| - |S_3| - |S_4| \notag \\ &+ |S_1 \cap S_2| + |S_1 \cap S_3| + |S_1 \cap S_4| + |S_2 \cap S_3| + |S_2 \cap S_4| + |S_3 \cap S_4| \notag \\ &- |S_1 \cap S_2 \cap S_3| - |S_1 \cap S_2 \cap S_4| - |S_1 \cap S_3 \cap S_4| - |S_2 \cap S_3 \cap S_4| \notag \\ &+ |S_1 \cap S_2 \cap S_3 \cap S_4| \notag \\ &= S_{r_1, r_2, r_3, r_4} - S_{l_1-1, r_2, r_3, r_4} - S_{r_1, l_2-1, r_3, r_4} - S_{r_1, r_2, l_3-1, r_4} - S_{r_1, r_2, r_3, l_4-1} \notag \\ &+ S_{l_1-1, l_2-1, r_3, r_4} + S_{l_1-1, r_2, l_3-1, r_4} + S_{l_1-1, r_2, r_3, l_4-1} + S_{r_1, l_2-1, l_3-1, r_4} + S_{r_1, l_2-1, r_3, l_4-1} + S_{r_1, r_2, l_3-1, l_4-1} \notag \\ &- S_{l_1-1, l_2-1, l_3-1, r_4} - S_{l_1-1, l_2-1, r_3, l_4-1} - S_{l_1-1, r_2, l_3-1, l_4-1} - S_{r_1, l_2-1, l_3-1, l_4-1} \notag \\ &+ S_{l_1-1, l_2-1, l_3-1, l_4-1} \notag \end{align}