尺取り法

Pythonのコードを交えて、尺取り法について説明します。

はじめに

先日、初めてAtCoderに参加*1し、尺取り法について理解したので簡単にまとめてみます。

尺取り法は、右端と左端の2つのインデックスを保持し、条件に応じて片方のインデックスを動かすことで、計算ループの最適化を図るアルゴリズムです。全探索の場合計算量は $O(N^2)$ ですが、尺取り法では $O(N)$ ですみます。

実例があったほうが分かりやすいので、例とともに紹介します。

ABC 032 C - 列

問題

長さ $N$ の非負整数列 $S=s_1,s_2,…,s_N$ と整数 $K$ があります。あなたの仕事は、以下の条件を満たす $S$ の連続する部分列のうち、最も長いものの長さを求めることです。部分列の長さは $1$ 以上の列でないといけません。

その部分列に含まれる全ての要素の値の積は、 $K$ 以下である。もし条件を満たす部分列が一つも存在しないときは、 $0$ を出力してください。

出典: C: 列 - AtCoder Beginner Contest 032 | AtCoder

解法

尺取り法で解く手順は以下のような感じです。

左端と右端に初期値を設定する
右端を進めたたとき、積が $K$ 以下の場合、右端を進める（図の青色）
右端が進められない場合、左端を進める（図の橙色）

このとき、右端と左端が同じ位置なら、右端も同時に進める（図の緑色）

2.に戻る

入力が $S = 4, 1, 7, 2, K=6$ のとき、イメージは以下のような感じです。

f:id:xkumiyu:20180527163851p:plain:w600

コード

n, k = map(int, input().split())
s = [int(input()) for _ in range(n)]

# 0が含まれる場合
if 0 in s: 
    ans = n
else:
    r, l = 0, 0 # 右端と左端のインデックス
    sum_s = 1 # 右端と左端に含まれる合計値
    ans = 0
    while r < n:
        # 右端を進めたときの積がK以下なら、右端を進める
        if sum_s * s[r] <= k:
            sum_s *= s[r]
            r += 1
            ans = max(ans, r - l)
        # 右端と左端を進める
        elif r == l:
            r += 1
            l += 1
        # 左端を進める
        else:
            sum_s //= s[l] # 左端を進めたときに取り除かれる要素を合計値からも取り除く
            l += 1
        # print('l={}, r={}, size={}, {}'.format(l, r, r - l, s[l:r]))
print(ans)

デバッグのコメントアウト部分を外すと以下のような出力が得られ、 $l=0, r=2$ のとき、つまり![ 4, 1] のとき部分列が最大となります。*2

l=0, r=1, size=1, [4]
l=0, r=2, size=2, [4, 1]
l=1, r=2, size=1, [1]
l=2, r=2, size=0, []
l=3, r=3, size=0, []
l=3, r=4, size=1, [2]

参考

【累積和、しゃくとり法】初級者でも解るアルゴリズム図解 - paiza開発日誌

www.slideshare.net

abc032 from AtCoder Inc.

その他の例題 ARC 022 B - 細長いお菓子

問題

高橋君は細長いお菓子を持っています。このお菓子は $N$ cm の長さのお菓子で、 $1$ cm ごとにブロックに分かれています。それぞれのブロックには $105$ 種類の味うちのいずれかの味がついていて、左端から $i$ 番目のブロックには $A_i$ 番目の味がついています。

高橋君はこのお菓子から、出来るだけ長い「同じ味のブロックが 2 つ以上含まれない、ひと繋がりになった部分」を切り出したいと思っています。最大で何 cm の部分を切り出すことが出来るでしょうか。ただし、切る場所はブロックとブロックの境界の部分のみとします。

出典: B: 細長いお菓子 - AtCoder Regular Contest 022 | AtCoder

解法

同じように、尺取り法を用いて解くことができます。ABC 032 C - 列では、条件が要素の積だったのが、今回は要素がすべて異なることになります。

この条件のチェックに、はじめは、

len(a[l:r + 1]) == len(set(a[l:r + 1]))

のように、重複があるかどうか確認していましたが、テストケース4つほどタイムアウトになり、部分点99点でした。

以下のコードのように、その値が使用されているかどうかを確認するリストを用意することで、この部分の時間が短縮でき、すべてのテストケースでパスできました。

コード

n = int(input())
a = [int(x) for x in input().split()]

l, r = 0, 0
v = [False for _ in range(max(a) + 1)] # 使用したどうかの確認するリスト
ans = 0
while r < n:
    # 右端を進める
    if not v[a[r]]:
        v[a[r]] = True # 使用したのでTrueにする
        r += 1
        ans = max(ans, r - l)
    # 右端と左端を進める
    elif l == r:
        l += 1
        r += 1
    # 左端を進める
    else:
        v[a[l]] = False # 左端を取り除く（Falseにする）
        l += 1
    # print('l={}, r={}, size={}, {}'.format(l, r, r - l, v))
print(ans)

ABC 098 D - Xor Sum 2

問題

長さ $N$ の整数列 $A$ があります。

次の条件を満たす整数 $l, r ( 1 \leq l \leq r \leq N )$ の組の個数を求めてください。

$A_l ,, xor ,, A_{l+1} ,, xor ,, \ldots ,, xor ,, A_r = A_l + A_{l+1} + \ldots + A_r$

出典: D: Xor Sum 2 - AtCoder Beginner Contest 098 | AtCoder

解法

条件のXORであることは、これまでの要素の和と新しい要素のANDが0であることを利用して、あとは同じように尺取り法で計算ループを進めます。

コード

n = int(input())
a = list(map(int, input().split()))

l, r = 0, 0
sum_a = 0
ans = 0
while r < n:
    # 右端を進める
    if sum_a & a[r] == 0:
        sum_a += a[r]
        r += 1
        ans += r - l
    # 右端と左端を進める
    elif l == r:
        l += 1
        r += 1
    # 左端を進める
    else:
        sum_a -= a[l]
        l += 1
    # print('l={}, r={}, size={}, ans={}, {}'.format(l, r, r - l, ans, a[l:r]))
print(ans)

*1:ABC 098に参加しましたが、D問題が解けませんでした。。

*2:図はインデックスが1からですが、コードでは0からであることに注意してください。