2020-07-19

Python 尺取法ライブラリ

プログラミングプログラミング-競技プログラミングプログラミング-競技プログラミング-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)

解説

$n$ ... リストのサイズを指定します。
$\mathrm{maximum}(A)$ ... 条件を満たすAの部分数列の最大の長さを返します。
$\mathrm{count}(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))

Python Imos法ライブラリ

プログラミングプログラミング-競技プログラミングプログラミング-競技プログラミング-Python

Imos法ライブラリ

imos法のPythonライブラリです。
imos法はある $[l, r)$ の区間に対して処理を行うクエリがたくさん与えられるときに使うアルゴリズムです。詳しくはいもすさんによる解説記事を見てください。

from itertools import accumulate
class Imos:
    def __init__(self, n):
        self.A = [0]*(n+1)
        self.n = n

    def __call__(self):
        dist = list(accumulate(self.A))[:self.n]
        self.A = [0] * (self.n+1)
        return dist

    def query(self, l, r):
        self.A[l] += 1
        self.A[r] -= 1

解説

$n$ ... リストのサイズを指定します。
$\mathrm{query}(l, r)$ ... 区間 $[l, r)$ に対するクエリを入力に取ります。
$\mathrm{call}$ ... クエリの反映を行いリスト状態を返します。

使用例

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

AtCoder ABC 035 C - オセロ

from itertools import accumulate
class Imos:
    # Folding

n, q, *LR = map(int, open(0).read().split())
imos = Imos(n)
for l, r in zip(LR[::2], LR[1::2]):
    imos.query(l-1, r)
print(*(x%2 for x in imos()), sep='')

AtCoder ABC 014 C - AtColor

from itertools import accumulate
class Imos:
    # Folding

n, *AB = map(int, open(0).read().split())
imos = Imos(1_000_001)
for l, r in zip(AB[::2], AB[1::2]):
    imos.query(l, r+1)
print(max(imos()))

AtCoder 東京海上日動プログラミングコンテスト2020 C - Lamps

# Python3 is TLE! Execute by Pypy3
from itertools import accumulate
class Imos:
    # Folding

n, k, *A = map(int, open(0).read().split())
imos = Imos(n)
for _ in range(min(100, k)):
    for i in range(n):
        imos.query(max(0, i-A[i]), min(i+A[i]+1, n))
    A = imos()
print(*A)

2020-07-15

Python 二部探索用ライブラリ

プログラミングプログラミング-競技プログラミングプログラミング-競技プログラミング-Python

二部探索ライブラリ

二部探索をしてくれるPythonライブラリです。
Pythonのbisectモジュールはソートされたリストにしか対応しておらず、単調増加関数に対する二部探索は自分で書く必要があります。以前のものより小回りが利くようにクラスモジュールで定義しました。

class BisectSearch:
    def __init__(self, f, l=0, r=10**9):
        self.f = f
        self.l = l
        self.r = r
    def __call__(self, dist):
        f = self.f
        l = self.l
        r = self.r
        if dist <= f(l):
            return l
        if f(r) <= dist:
            return r
        while r-l > 1:
            n = (r+l) // 2
            if f(n) <= dist:
                l = n
            else:
                r = n
        return l

入力

$l, r$ ... 閉区間 $[l, r$ ] で探索したい範囲を指定します。
$f$ ... $[l, r$ ] $\mapsto \mathbb{R}$ となる単調増加関数 $f$ を指定します。
$\mathrm{dist}$ ... 探索したい値を指定します。

出力

$f(t) \le dist \lt f(t+1)$ を満たすときの $t$ を返します。ただし $dist \lt f(l), f(r) \le dist$ のときはそれぞれ $l, r$ を返します。

使用例

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

AtCoder ABC 146 C - Buy an Integer

class BisectSearch:
    # Folding

a, b, x = map(int, input().split())
def f(n):
    return a*n + b*len(str(n))
print(BisectSearch(f, l=0, r=10**9)(x))

AtCoder CODE FESTIVAL 2014 決勝 - C - N進数

class BisectSearch:
    # Folding

a = int(input())
def f(n):
    return sum(int(s) * n**i for i, s in enumerate(str(n)[::-1]))
k = BisectSearch(f, 1, 10**16)(a)
print(k if f(k) == a and k >= 10 else -1)

AtCoder CODE FESTIVAL 2016 Final B - Exactly N points

class BisectSearch:
    # Folding

n = int(input())
def f(n):
    return (n+1)*n // 2
m = BisectSearch(f, l=1, r=10**7)(n)
m += f(m) != n
A = set(range(1, m+1))
A.discard(f(m) - n)
print(*A, sep='\n')

2020-05-10

Python 階乗・順列・二項係数を素数を割った余りを求めるライブラリ

プログラミングプログラミング-競技プログラミングプログラミング-競技プログラミング-Python

階乗・順列・二項係数の剰余を求めるライブラリ

競プロでよく使うやつをライブラリ化しました。Moduloが素数でない時は上手く動きません。

class Factorial():
    def __init__(self, mod=10**9 + 7):
        self.mod = mod
        self._factorial = [1]
        self._size = 1
        self._factorial_inv = [1]
        self._size_inv = 1
    
    def __call__(self, n):
        '''n! % mod '''
        return self.fact(n)
    
    def fact(self, n):
        '''n! % mod '''
        if n >= self.mod:
            return 0
        self.make(n)
        return self._factorial[n]
    
    def fact_inv(self, n):
        '''n!^-1 % mod '''
        if n >= self.mod:
            raise ValueError('Modinv is not exist! arg={}'.format(n))
        self.make_inv(n)
        return self._factorial_inv[n]
    
    def comb(self, n, r):
        ''' nCr % mod '''
        if r > n:
            return 0
        t = self.fact_inv(n-r)*self.fact_inv(r) % self.mod
        return self(n)*t % self.mod
    
    def comb_with_repetition(self, n, r):
        ''' nHr % mod '''
        t = self.fact_inv(n-1)*self.fact_inv(r) % self.mod
        return self(n+r-1)*t % self.mod
    
    def perm(self, n, r):
        ''' nPr % mod '''
        if r > n:
            return 0
        return self(n)*self.fact_inv(n-r) % self.mod
    
    @staticmethod
    def xgcd(a, b):
        ''' return (g, x, y) such that a*x + b*y = g = gcd(a, b) '''
        x0, x1, y0, y1 = 0, 1, 1, 0
        while a != 0:
            (q, a), b = divmod(b, a), a
            y0, y1 = y1, y0 - q * y1
            x0, x1 = x1, x0 - q * x1
        return b, x0, y0
    
    def modinv(self, n):
        g, x, _ = self.xgcd(n, self.mod)
        if g != 1:
            raise ValueError('Modinv is not exist! arg={}'.format(n))
        return x % self.mod
    
    def make(self, n):
        if n >= self.mod:
            n = self.mod
        if self._size < n+1:
            for i in range(self._size, n+1):
                self._factorial.append(self._factorial[i-1]*i % self.mod)
            self._size = n+1
    
    def make_inv(self, n):
        if n >= self.mod:
            n = self.mod
        self.make(n)
        if self._size_inv < n+1:
            for i in range(self._size_inv, n+1):
                self._factorial_inv.append(self.modinv(self._factorial[i]))
            self._size_inv = n+1

解説

${n!}$ を_factorial, ${}n!^{-1}$ を_factorial_invに記録するみたいな感じです。モジュロ逆数を求めるアルゴリズムをフェルマーの小定理に頼っているため、モジュロが素数でないと上手く動作しません。拡張ユークリッド互除法をつかって逆元を計算したほうが拡張性が高いですが、大目に見てください。
拡張ユークリッド互除法に変更しました！

他にも順列や二項係数、重複組合せなどをサポートしてます。

下のページの先頭クエリ100個で動作テストしました。
yukicoder.me

コードはこんな感じ。

class Factorial():
# 省略

fact = Factorial()
comb = fact.comb
perm = fact.perm
comb_with_repetition = fact.comb_with_repetition

inputs = open(0).read().split()
for i in range(int(inputs[0])):
    q, n, r = inputs[i+1].replace('(', ',').replace(')', '').split(',')
    n, r = map(int, (n, r))
    if q == 'C':
        print(comb(n, r))
    elif q == 'P':
        print(perm(n, r))
    else:
        print(comb_with_repetition(n, r))

その他

~~拡張ユークリッドを実装したら更新するかも。再帰だとPypyで遅くなるので非再帰で書きたいんですが、まあがんばります。~~
20200719 改善済み

2020-05-08

Welcome to JupyterLab in AtCoder

プログラミングプログラミング-競技プログラミングプログラミング-競技プログラミング-Python

JupyterLabとは

JupyterLabは、Project Jupyterで次世代のWebベースのユーザーインターフェイスです。Jupyter notebookの後継として開発され、データサイエンティストの界隈では有名なアプリケーションです。ざっくりとPythonのIDEみたいな使い方ができます。ちなみにAnacondaをインストールすればすぐに使えます。

どんな感じで使ってんの

”百聞は一見に如かず”ということで私のコーディングの雰囲気を見せます。

f:id:ybrkbr:20200508022649p:plain — Jupyter Labでのコーディングの見た目

解説

左上、右上

Jupyterで標準入力を楽にするためのシステムです。出力セルにTextareaを出し、そこに書いてあるテキストにopen(0)でアクセスできます。下にコピペ用コードを示しますが、別途インストールが必要なものがあるのでここを参考にインストールしてください。コードはこの記事を参考にしました。
JupyterLabでは出力セルを右クリックして「Create New View for Output」ってすると別ペインに出すことができて入力をインタラクティブにかえつつ実行できて便利です。

from ipywidgets import Textarea
import io

class Open():
    def __init__(self):
        self.text = ''

    def __call__(self, file, *args, **kwargs):
        if file == 0:
            return io.StringIO(self.text)
        return open_(file, *args, **kwargs)

    def updater(self, change):
        self.text = change["new"]

open = Open()
text_area = Textarea()
text_area.observe(open.updater, names='value')
display(text_area)

左下

コードを書いていく場所です。今回はAtCoderのpractice問題のA - Welcome to AtCoderをやってます。

右下

これはまあ正直なくてもいいんですが、「New Console for Notebook」を開いてwhosで変数の一覧を眺めたり、リストの中身を調べたりします。

その他

自分の中で環境が変わったり、質問などがきたら追記します。

2020-04-07

Python グリッド用ライブラリ

プログラミングプログラミング-競技プログラミングプログラミング-競技プログラミング-Python

グリッド用ライブラリ

競プロでグリッド系の問題が出てとき、手間取りがちなのでライブラリにまとめました。
numpy実装でnumpy特有のスライスも対応してます。
np.allやnp.whereなどndarrayを引数にする関数を使いたい場合は、Grid.gridをお使いください()

import numpy as np
class Grid():
    def __init__(self, grid, w=0, h=0, function=lambda x: x):
        self.w = w = w if w else len(grid[0])
        self.h = h = h if h else len(grid)
        dtype = type(function(grid[0][0]))
        self.grid = np.empty((h, w), dtype=dtype)
        for i, row in zip(range(h), grid):
            for j, val in zip(range(w), row):
                self.grid[i][j] = function(val)
    
    def is_valid_x(self, x):
        return 0 <= x < self.w
    def is_valid_y(self, y):
        return 0 <= y < self.h
    def is_valid_xy(self, x, y):
        return self.is_valid_x(x) and self.is_valid_y(y) 
    
    def __iter__(self):
        return iter(self.grid)
    def __repr__(self):
        return '\n'.join([' '.join(map(str, row)) for row in self.grid])
    def __getitem__(self, x):
        return self.grid[x]
    def __setitem__(self, x, val):
        self.grid[x] = val

def dfs(root):
    x, y = root
    grid[y, x] = 0
    stack = [root]
    while stack:
        x, y = stack.pop()
        for dx, dy in zip([1, 0, -1, 0], [0, 1, 0, -1]):
            if grid.is_valid_xy(x+dx, y+dy) and grid[y+dy, x+dx]:
                stack.append((x+dx, y+dy))
                grid[y, x] = 0

from collections import deque
def bfs(root):
    x, y = root
    grid[y, x] = 0
    queue = deque([root])
    while queue:
        x, y = queue.popleft()
        for dx, dy in zip([1, 0, -1, 0], [0, 1, 0, -1]):
            if grid.is_valid_xy(x+dx, y+dy) and grid[y+dy, x+dx]:
                queue.append((x+dx, y+dy))
                grid[y+dy, x+dx] = 0

解説

引数にはインデクシングが可能な2次元イテレータを入れればOKです。
要はAtCoderの入力の2次元リストや文字列のリストをそのまま入れれば大丈夫です。
入力がmapやgeneratorの場合は $w, h$ を引数にする必要があるので注意！
functionはmapみたいな感じで初期化時にgrid全体に関数が適用されます。
ついでにdfs,bfs実装のテンプレートも書いてます。

使用例

AtCoder ABC 96 C - Grid Repainting 2
AtCoder ABC 151 D - Maze Master
AtCoder AGC 43 A - Range Flip Find Route

# C - Grid Repainting 2
from itertools import product
from collections import deque
import numpy as np
def Grid():
## 略
inputs = open(0).readlines()
h, w = map(int, inputs[0].split())
grid = Grid(inputs[1:], function=lambda x: int(x == '#'))

def dfs(root):
    count = 1
    x, y = root
    grid[y, x] = 0
    stack = [root]
    while stack:
        x, y = stack.pop()
        grid[y, x] = 0
        for dx, dy in zip([1, 0, -1, 0], [0, 1, 0, -1]):
            if grid.is_valid_xy(x+dx, y+dy) and grid[y+dy, x+dx]:
                stack.append((x+dx, y+dy))
                count += 1
    return count
    
ans = 'Yes'
for i, j in product(range(h), range(w)):
    if grid[i, j]:
        if dfs((j, i)) == 1:
            ans = 'No'
            break
print(ans)

# D - Maze Master
from copy import deepcopy
from collections import deque
from itertools import product
import numpy as np
def Grid():
## 略
inputs = open(0).readlines()
h, w = map(int, inputs[0].split())
grid_origin = Grid(inputs[1:], function=lambda x: int(x == '.'))

def bfs(root):
    x, y, _ = root
    grid[y, x] = 0
    queue = deque([root])
    while queue:
        x, y, d = queue.popleft()
        for dx, dy in zip([1, 0, -1, 0], [0, 1, 0, -1]):
            if grid.is_valid_xy(x+dx, y+dy) and grid[y+dy, x+dx]:
                queue.append((x+dx, y+dy, d+1))
                grid[y+dy, x+dx] = 0
    return d
    
ans = 0
for i, j in product(range(h), range(w)):
    if grid_origin[i, j]:
        grid = deepcopy(grid_origin)
        ans = max(ans, bfs((j, i, 0)))
print(ans)

# Grid Compression
import numpy as np
def Grid():
## 略
inputs = [s.strip() for s in open(0).readlines()]
h, w = map(int, inputs[0].split())
grid = Grid(inputs[1:])

dp = Grid(np.zeros((h, w), np.int))
dp[0, 0] = int(grid[0, 0] == '#')
for i in range(1, w):
    dp[0, i] = dp[0, i-1] + int(grid[0, i-1] + grid[0, i] == '.#')
for i in range(1, h):
    dp[i, 0] = dp[i-1, 0] + int(grid[i-1, 0] + grid[i, 0] == '.#')
x = y = 1
while not (x == w-1 and y == h-1):
    for i in range(x, w):
        dp[y, i] = min(dp[y, i-1] + int(grid[y, i-1] + grid[y, i] == '.#'),
                       dp[y-1, i] + int(grid[y-1, i] + grid[y, i] == '.#'))
    for i in range(y+1, h):
        dp[i, x] = min(dp[i-1, x] + int(grid[i-1, x] + grid[i, x] == '.#'),
                       dp[i, x-1] + int(grid[i, x-1] + grid[i, x] == '.#'))
    x, y = min(x+1, w-1), min(y+1, h-1)
print(min(dp[y-1, x] + int(grid[y-1, x] + grid[y, x] == '.#'), 
          dp[y, x-1] + int(grid[y, x-1] + grid[y, x] == '.#')))

2020-02-19

旧 Python 二部探索デコレータ関数

プログラミングプログラミング-競技プログラミングプログラミング-競技プログラミング-Python

追記（2020/07/15）

二部探索用コードはアップデートされました。 ⇨更新先サイト

二部探索デコレータ関数

二部探索用のPythonコードのメモ書きです。デコレータを使っていい感じにしてます。

def bisect_search(l=0, r=10**6):
    def _bisect_search(f, l=l, r=r):
        if not f(l):
            return l
        if f(r):
            return r
        while r-l > 1:
            n = (r+l) // 2
            if f(n):
                l = n
            else:
                r = n
        return l
    return _bisect_search

解説

$l, r$ ... 閉区間で探索したい範囲を指定します。
$f$ ... intからboolの写像です。定義域のある値を境に値域がTrue⇒Falseになるような単調な写像である必要があります。

使用例

現状少ないので情報をくれればアップデートします。

AtCoder ABC 146 C - Buy an Integer

from math import log10
inputs = open(0).read().split('\n')
 
def bisect_search(l=0, r=10**6):
    def _bisect_search(f, l=l, r=r):
        if not f(l):
            return l
        if f(r):
            return r
        while r-l > 1:
            n = (r+l) // 2
            if f(n):
                l = n
            else:
                r = n
        return l
    return _bisect_search
  
a, b, x = map(int, inputs[0].split())
@bisect_search(l=0, r=10**9)
def f(n):
    return n == 0 or a*n + b*int(log10(n)+1) <= x
print(f)