Python 尺取法ライブラリ

尺取法ライブラリ

尺取法のPythonライブラリです。英語でTwo Pointersっていうらしいです。 数列のうち、ある条件を満たす部分数列を数え上げたり、最大の長さの部分数列を求めたりすることができます。

from copy import deepcopy
from operator import add, sub
class TwoPointers:
    def __init__(self, cond, init=0, right=add, left=sub):
        self.cond = cond
        self.init = init
        self.next_right = right
        self.next_left = left
    
    def __call__(self, A, uniq=True):
        cond = self.cond
        next_right, next_left = self.next_right, self.next_left
        s = deepcopy(self.init)
        n = len(A)
        X = {}
        r = 0
        for l in range(n):
            if not uniq:
                X[l] = r
            while r < n and cond(s, A[r]):
                s = next_right(s, A[r])
                r += 1
                X[l] = r
            if l == r:
                r += 1
            else:
                s = next_left(s, A[l]) if l+1 != r else deepcopy(self.init)
        return list(X.items())
    
    def maximum(self, A):
        X = self.__call__(A)
        return max(r-l for l, r in X) if X else 0
    
    def count(self, A):
        X = self.__call__(A, False)
        return sum(r-l for l, r in X)

解説

... リストのサイズを指定します。
... 条件を満たすAの部分数列の最大の長さを返します。
... Aの部分数列で条件を満たすものを数え上げます。

使用例

使用したコードがあれば適宜アップデートしていきます。

AtCoder ABC 032 C - 列

from copy import deepcopy
from operator import add, sub, mul, floordiv
class TwoPointers:
    # Folding

n, k, *S = map(int, open(0).read().split())
if 0 in S:
    print(n)
else:
    tp = TwoPointers(lambda s, n: s*n <= k, 1, mul, floordiv)
    print(tp.maximum(S))

AtCoder ABC 098 D - Xor Sum 2

from copy import deepcopy
from operator import add, sub
class TwoPointers:
    # Folding

n, *A = map(int, open(0).read().split())
print(TwoPointers(lambda s, n: s + n == s ^ n).count(A))

AtCoder ABC 038 C - 単調増加

from copy import deepcopy
from operator import add, sub
class TwoPointers:
    # Folding

n, *A = map(int, open(0).read().split())
tp = TwoPointers(lambda s, n: s < n, 0, lambda s, n: n, lambda s, n: s)
print(tp.count(A))

AtCoder ARC 022 B - 細長いお菓子

from copy import deepcopy
from operator import add, sub
class TwoPointers:
    # Folding
n, *A = map(int, open(0).read().split())
def next_right(s, n):
    s.add(n)
    return s
def next_left(s, n):
    s.remove(n)
    return s
tp = TwoPointers(lambda s, n: n not in s, set(), next_right, next_left)
print(tp.maximum(A))