見出し画像
Photo by takahashi8

二次元累積和は「左上隅から」がポイント (ABC331 D - Tile Pattern)

こじま

はじめに

大和証券プログラミングコンテスト2023（AtCoder Beginner Contest 331）に参加しました。コンテストの復習で、二次元累積和に初めて触れました。二次元領域内の和が高速に計算できて、しかも複雑な場合分けも不要でとても気持ちよかったです！

どうやら二次元累積和を上手に使うためには、ちょっとしたポイントがあるようです。この記事では、競技プログラミングで二次元累積和を活用するための（自分なりの）コツを、実際の問題を通して解説します。


問題

平面に同じ模様のタイルが敷き詰められています。1枚のタイルにはN行N列のマス目があり、マス目は白か黒に塗られています。

「(A,B)を左上隅、(C,D)を右下隅とする長方形領域に含まれる黒マスの個数を答えよ」というクエリに答えてください。

制約

  • 1 ≤ N ≤ 1000

  • 1 ≤ Q ≤ 2×10^5

  • 0 ≤ A ≤ C < 10^9

  • 0 ≤ B ≤ D < 10^9

  • N, Q, A, B, C, D は全て整数

解説

下準備をして、クエリを O(1)で処理することを目指します。
まずは二次元累積和を作りましょう。

二次元累積和を作る

N, Q =　map(int,input().split())
P = [input() for _ in range(N)]

Cum = [[0]*(N+1)]
for p in P :
    tmp = [0]
    for ch in p :
        tmp.append(tmp[-1]+(ch=="B"))
    for i, k in enumerate(Cum[-1]) :
        tmp[i] += k
    Cum.append(tmp)

各行で横方向に累積和を取ってから、上の行を加算しています。
直感的で分かりやすいと思ってます。ここは好みの方法で書きましょう。

簡単な問題を考える

最初から任意のケースを考えるとややこしくなりがちです。まずは問題を簡略化して、実験から始めてみましょう。具体的には、一枚のタイルの上だけに限定して考えてみます。

制約を強めて、

  • 0 ≤ A ≤ C < N

  • 0 ≤ B ≤ D < N

とすることで、一枚のタイル上の問題となります。

左上隅から(h,w)までの計算は、二次元累積和そのままです。

def g(h,w) :
    return Cum[h][w]

(A,B)から(C,D)までの長方形領域も、二次元累積和の性質から求めることができます。

def f(A,B,C,D) :
    return g(C,D) - g(A,D) - g(C,B) + g(A,B)

関数 f(A,B,C,D) は右半開領域を計算しています。
なのでクエリでは、(A,B)から(C+1,D+1)を答えることにしましょう。

for _ in range(Q) :
    a, b, c, d = map(int,input().split())
    print(f(a,b,c+1,d+1))

これで、一枚のタイルの上に限定した問題に答えることができました。ここまでが二次元累積和の基本形ですね。

任意の領域を計算するには？

さて、一枚のタイルの上の領域は計算できるようになりました。それでは、複数のタイルをまたぐ領域を計算するためにはどうしたらよいでしょう？

一枚のタイルの上であれば任意の領域が計算できるのですから、タイルごとにクエリの領域を分割することで解決できそうです。けれども、この方法は実装がややこしくなります。上下左右にはみ出している部分をもれなく計算するのは骨が折れます。実際に試したときも、時間がかかった上にバグらせました。

ここは公式解説の通りに、g(h,w)を拡張します。一枚のタイル上だけでなく全域で、左上隅から(h,w)までの計算ができるようにしましょう。

def g(h,w) :
    if h <= N and w <= N : return Cum[h][w]
    hq, hr = h//N, h%N
    wq, wr = w//N, w%N
    rtn = 0
    rtn += g(N,N)*hq*wq
    rtn += g(hr,N)*wq
    rtn += g(N,wr)*hq
    rtn += g(hr,wr)
    return rtn

左上隅からの領域なので、タイルからはみ出す部分は右と下だけです。おかげで計算がシンプルですね。再帰を使うとなんとなくカッコいいですが、愚直に累積和Cumを使ってもよいと思います。

def g(h,w) :
    hq, hr = h//N, h%N
    wq, wr = w//N, w%N
    return Cum[N][N]*hq*wq + Cum[hr][N]*wq + Cum[N][wr]*hq + Cum[hr][wr]

これで g(h,w) が全域に拡張できました。f(A,B,C,D)もそのまま全域で使うことができるので、以下は簡単な問題のときのコードをそのまま使うことができます。

def f(A,B,C,D) :
    return g(C,D) - g(A,D) - g(C,B) + g(A,B)

for _ in range(Q) :
    a, b, c, d = map(int,input().split())
    print(f(a,b,c+1,d+1))

これで任意の領域をO(1)で計算して、クエリに答えられるようになりました。やった！

まとめ

ABC331-Dは、二次元累積和の応用問題でした。二次元累積和では「左上隅から」の計算を元にして、任意の領域を高速に求めています。累積和を取った行列ではなく「左上隅から」の計算こそが要です。

「左上隅から」の計算範囲が制限されている問題でも、それを全域に拡張する方法を考えてみることで、シンプルな解答が得られるかもしれません。

参照

公式解説です。図もあるので計算をイメージしやすいです。

自分の最終的なACコードです。23行でした。

「累積和とは何か」から始まって発展的トピックまでまとめられた、お化けのような記事。すごい。二次元累積和の例題も2問扱っています。

いいなと思ったら応援しよう！