Pythonでmultisetの代用ができるデータ構造

ABC170のE問題でmultisetを使う問題が出題されましたが、Pythonにはmultisetがありません（さらに順序付きの普通の集合もないですが...）。そこで以下の処理を行うことができるデータ構造、HeapDictを作りました（ソースコードも一番下にあります）。

〇O(logn)で要素の挿入
〇O(logn)で要素の削除
〇O(1)で要素の存在確認
〇O(1)で最小値の取得

これは単純にheapqの上位互換です。C++のmultisetやsetとかなり似ています。二分探索ができないだけです。

PythonのsetやdictはO(1)で要素の挿入および削除ができますが、順序付きでないため最小値の取得にはO(n)かかります。また、PythonのheapqはO(logn)で要素の挿入ができ、O(1)で最小値の取得ができますが、最小値以外の削除ができない？はずです。しかし、これらを組み合わせることでmultiset的な挙動をするデータ構造を作ることができます。

具体的な実装方法を述べます。
heapqとdictを用意します。

def __init__(self):
    self.h=[]
    self.d=dict()

要素の挿入はどちらもできるため、普通に挿入します。計算量はdictでO(1)、heapqでO(logn)となります。

def insert(self,x):
    heapq.heappush(self.h,x)
    if x not in self.d:
        self.d[x]=1
    else:
        self.d[x]+=1

要素の削除はheapqでは普通はできませんが、削除を遅延させることで解決します。このデータ構造で取得したいものは最小値であるため、最小値さえ矛盾していなければ要素の削除は遅延させてもかまわないということになります。遅延方法としては、まずdictで普通に要素の削除をします。次にheapqの方の最小値を確認します。この最小値がdictにすでに存在しないとなれば削除し、最小値がdictに存在するようになるまで削除し続けます。遅延削除の回数は多くても通常の削除の回数と同じとなります。

def erase(self,x):
    if x not in self.d or self.d[x]==0:
        print(x,"is not in HeapDict")
        exit()
    else:
        self.d[x]-=1

    while len(self.h)!=0:
        if self.d[self.h[0]]==0:
            heapq.heappop(self.h)
        else:
            break

要素の存在確認はdictの方を見ればO(1)で行えます。

def is_exist(self,x):
    if x in self.d and self.d[x]!=0:
        return True
    else:
        return False

最小値の取得は、最小値が矛盾していないという保証があるため、heapqの方を見ればO(1)で行えます。

def get_min(self):
    return self.h[0]

最後にこれらをまとめたものがこちらになります。

import heapq
from sys import exit
class HeapDict:
    def __init__(self):
        self.h=[]
        self.d=dict()

    def insert(self,x):
        heapq.heappush(self.h,x)
        if x not in self.d:
            self.d[x]=1
        else:
            self.d[x]+=1

    def erase(self,x):
        if x not in self.d or self.d[x]==0:
            print(x,"is not in HeapDict")
            exit()
        else:
            self.d[x]-=1

        while len(self.h)!=0:
            if self.d[self.h[0]]==0:
                heapq.heappop(self.h)
            else:
                break

    def is_exist(self,x):
        if x in self.d and self.d[x]!=0:
            return True
        else:
            return False

    def get_min(self):
        return self.h[0]

使い方としてはPython標準のsetっぽく、insert(値)としたり、erase(値)としたりすればよいです。

競プロでよく使うPythonモジュール

備忘録
〇sys
・sys.stdin.readlineは入力の高速化に必須
・sys.exitもたまに使う

〇math
・gcdをよく使う
・factorialは階乗を返す
・三角関数などもここにある

〇decimal
・精度の高い小数を扱える

〇collections
・dequeをよく使う
・Counterも多少使う

〇heapq
・優先度付きキュー

〇bisect
・二分探索

〇itertools
・accumulateで累積和を計算できる
・permutationsで順列を生成する
・combinationsで組み合わせを生成する
・combinations_with_replacement重複ありの組み合わせを生成する

〇numpy
・標準ライブラリではないが、高速化に大きく貢献する

〇scipy
・numpyと同様

numpy用練習問題

numpyで簡潔に書けたり劇的に高速化できたりする問題一覧
競プロでPythonを使っているが、numpyをあまり使っていないという人はやってみるとnumpyの効果を実感できると思います。

D - Knapsack 1
典型的なナップサックDP。Pythonで愚直にfor文を回すやり方では間に合わないが、numpyを使ってfor文を減らすと劇的に高速化できる例。また、numpy.maximumで２つの配列の要素を比較していき、大きい方の要素を選んだ配列を作ることができる。
Submission #14017275 - Educational DP Contest

D - Lamp
こちらもKnapsack 1と同様に愚直にやると間に合わないが、工夫して配列ごと計算することでfor文を減らすことができる。
Submission #14016436 - AtCoder Beginner Contest 129

D - Good Grid
こちらは愚直にやると前処理がボトルネックとなりTLEしてしまう問題だが、簡単にfor文を減らすことができる。
Submission #14023232 - AtCoder Beginner Contest 099

D - AtCoder Express 2
numpy.ndarray.sum（総和）とnumpy.ndarray.cumsum（累積和）が使える。通常のリストでは列方向にsumを取ることができないが、numpyを使えばそれが可能となる。
Submission #13995334 - AtCoder Beginner Contest 106

B - Ruined Square
numpy.dotで行列の積を利用することができる。
Submission #14009262 - AtCoder Beginner Contest 108

D - Patisserie ABC
こちらも行列の積を利用する解法がある。
Submission #14001450 - AtCoder Beginner Contest 100

C - Sum of gcd of Tuples (Easy)
これは3重for文で愚直にやっても間に合うが、numpy.gcd.outerで簡潔に書ける。
Submission #14020240 - AtCoder Beginner Contest 162

C - 億マス計算
D - Pairs
二分探索を使う問題。Python標準のbisectを使ってもTLEしてしまう厳しい問題だが、numpy.searchsortedなどを駆使すると簡潔に書ける上、劇的に高速化できる。
Submission #13602083 - AtCoder Regular Contest 037
Submission #13609335 - AtCoder Beginner Contest 155

（AtCoder）ABC169

A
A*Bを出力する。

B
単純に計算しながら結果が10^18を超えたところで打ち切りにするやり方では0があるケースに対応できない。そのため、まず0があるかどうか確認する。

C
PythonならDecimaiでゴリ押せる。

D
素因数分解してそれぞれの素因数pについてp^1,p^2,p^3,...で割れるなら割っていく。

E,Fはわかりませんでした。
E
偶奇で分けるぐらいしかわからず

F
DPをするのかな？くらいしかわからず

（AtCoder）numpyの高速化テクニック

AtCoderのPython3はnumpyを使うことができる。その高速化テクニックをまとめる。
時間計測はAtCoderのコードテストを使用。
※追記中

〇基本的な考え方

クリックで展開

・for文を減らす。→numpyの配列を使ってまとめて計算する。
＜例題＞
D - Lamp
D - Knapsack 1
C - Sum of gcd of Tuples (Easy)

〇二分探索

クリックで展開

N=10**6
大きさNのリストでN回二分探索をする。計算量O(NlogN)。
numpyなし：bisectを使用
numpyあり：numpy.searchsortedを使用

import time
import bisect
import numpy as np
n=10**6
a=list(range(n))
b=list(range(n))
c=np.array(range(n),dtype="int64")
d=np.array(range(n),dtype="int64")

#numpyなし
t1=time.time()
for i in range(n):
    j=bisect.bisect_left(a,b[i])
t2=time.time()
print(t2-t1)

#numpyあり
t3=time.time()
js=np.searchsorted(c,d)
t4=time.time()
print(t4-t3)

結果
numpyなし：約0.6s弱
numpyあり：約0.08s
大幅な高速化に成功。

＜例題＞
C - 億マス計算
D - Pairs

〇条件を満たす値を調べる

クリックで展開

N=10**6
大きさNのリストのすべての要素に対し、ある値xとの大小関係を調べる。計算量O(N)。
numpyなし1：for文を回して1つずつ比較
numpyなし2：リスト内包表記を使用
numpyあり：b[b>x]を使用

import time
import numpy as np
n=10**6
a=list(range(n))
b=np.array(range(n),dtype="int64")

x=100
#numpyなし1
t1=time.time()
cnt=0
for i in range(len(a)):
    if a[i]>x:
        cnt+=1
print(cnt)
t2=time.time()
print(t2-t1)

#numpyなし2
t3=time.time()
print(len([a[i] for i in range(n) if a[i]>x]))
t4=time.time()
print(t4-t3)

#numpyあり
t5=time.time()
print(len(b[b>x]))
t6=time.time()
print(t6-t5)

結果
numpyなし1：約0.18s
numpyなし2：約0.12s
numpyあり：約0.005s
大幅な高速化に成功。

＜例題＞
D - Pairs

〇累積和

クリックで展開

N=10**6
大きさNのリストの累積和を求める。計算量O(N)。
numpyなし：itertoolsのaccumulateを使用
numpyあり：cumsumを使用

import time
from itertools import accumulate
import numpy as np
n=10**6
a=list(range(n))
b=np.array(range(n),dtype="int64")

t1=time.time()
s1=list(accumulate(a))
t2=time.time()
print(t2-t1)

t3=time.time()
s2=b.cumsum()
t4=time.time()
print(t4-t3)

結果
numpyなし：約0.04s
numpyあり：約0.004s
それなりに速くなる。

＜例題＞
D - AtCoder Express 2

〇総和

クリックで展開

N=10**6
大きさNのリストの総和を求める。計算量O(N)。
numpyなし：sum(a)を使用
numpyあり：b.sum()を使用

import time
import numpy as np

a=list(range(10**6))
b=np.array(range(10**6),dtype="int64")

#numpyなし
t1=time.time()
print(sum(a))
t2=time.time()

#numpyあり
t3=time.time()
print(b.sum())
t4=time.time()

print(t2-t1)
print(t4-t3)

結果
numpyなし：約0.007s
numpyあり：約0.001s
それほど変化なし。

＜例題＞
D - AtCoder Express 2

〇最大最小

クリックで展開

N=10**6
大きさNのリストの最大値（最小値）を求める。計算量O(N)。
numpyなし：max(a)を使用
numpyあり：b.max()を使用

import time
import numpy as np

a=list(range(10**6))
b=np.array(range(10**6),dtype="int64")

#numpyなし
t1=time.time()
print(max(a))
t2=time.time()

#numpyあり
t3=time.time()
print(b.max())
t4=time.time()

print(t2-t1)
print(t4-t3)

結果
numpyなし：約0.018s
numpyあり：約0.0012s
それなりに速くなる。

※注意
for文の繰り返し回数が多く、min(max)を使用する回数が多いが、配列の大きさは小さいという場合にはnumpyのmin(max)とリストのmin(max)の速度が逆転することが分かった。つまり、numpyのmin(max)を繰り返し使用するときには遅くなる可能性がある。

import time
import numpy as np
def main():
    a=np.array(range(10**5),dtype=np.int64)
    b=list(range(10**5))
    c=np.array(range(10**2),dtype=np.int64)
    d=list(range(10**2))

    t1=time.time()
    for i in range(10**2):
        x=a.min()
    t2=time.time()

    t3=time.time()
    for i in range(10**2):
        x=min(b)
    t4=time.time()

    t5=time.time()
    for i in range(10**5):
        x=c.min()
    t6=time.time()

    t7=time.time()
    for i in range(10**5):
        x=min(d)
    t8=time.time()

    print(t2-t1)
    print(t4-t3)
    print(t6-t5)
    print(t8-t7)

if __name__=="__main__":
    main()

結果
〇for文少、配列大の場合
numpyあり：約0.01s
numpyなし：約0.16s
〇for文多、配列小の場合
numpyあり：約0.5s
numpyなし：約0.17s

〇2つの配列の要素を比較し、大きい方（小さい方）を選んでいく

クリックで展開

N=10**6
numpyなし1：for文を回して1つずつ比較
numpyなし2：リスト内包表記を使用
numpyあり：numpy.maximumを使用

import time
import numpy as np

n=10**6
x=100
a1=list(range(n))
a2=list(range(n))
b1=[x]*n
b2=[x]*n
c=np.array(range(n),dtype="int64")
d=np.full(n,x,dtype="int64")

#numpyなし1
t1=time.time()
for i in range(n):
    if a1[i]<b1[i]:
        a1[i]=b1[i]
t2=time.time()

#numpyなし2
t3=time.time()
a2=[b2[i] if a2[i]<b2[i] else a2[i] for i in range(n)]
t4=time.time()

#numpyあり
t5=time.time()
c=np.maximum(c,d)
t6=time.time()

print(t2-t1)
print(t4-t3)
print(t6-t5)

結果
numpyなし1：約0.13s
numpyなし2：約0.14s
numpyあり：約0.003s
大幅な高速化に成功。一般にリスト内包表記の方が速いと言われるが、今回はリスト内包表記の方がなぜか遅い。

＜例題＞
D - Knapsack 1

ABC134 E - Sequence Decomposing

atcoder.jp
解法
貪欲的に考える。
例えば2 1 4 2のような場合、2と1を色aと色bで塗ったとする（これらは異なる色で塗るしかない）。そうすると、次の4は2と同じ色aで塗った方がよいことが分かる。なぜなら、
〇色aの最大値が4、色bの最大値が1
〇色aの最大値が2、色bの最大値が4
前者の方が「状況がよい」ため、つまり次の2を色bで塗れるためである。

より一般的に考えて整数nを塗るときは、色a、色b、色c、・・・のうち、その最大値が整数nより小さいものの中で最大となる色を選んで塗ればよく、整数nより小さな最大値が存在しないときは新しい色で塗ればよい。これを実装するには各色の最大値を要素に持つ配列を用意して二分探索を行えばよい（この配列は常に広義単調減少となるため常に二分探索が可能である）。整数nをすでに存在する色で塗るときにはその色の最大値を表す要素を整数nで更新し、整数nを新しい色で塗るときには配列の末尾に追加する、という操作を順番に行えばよい。

ソースコード（Python3）

from sys import stdin
def main():
    #入力
    readline=stdin.readline
    n=int(readline())
    a=[int(readline()) for _ in range(n)]

    li=[]
    for i in range(n):
        if i==0:
            li.append(a[i])
        else:
            #二分探索
            l=-1
            r=len(li)
            while l<r-1:
                m=(l+r)//2
                if li[m]>=a[i]:
                    l=m
                else:
                    r=m
            if r==len(li):
                li.append(a[i])
            else:
                li[r]=a[i]

    print(len(li))
    
if __name__=="__main__":
    main()

（Python3）競プロの入力

Python3でAtCoderなど競プロをやる人向け
入力を受け取る方法についてのまとめです。これらを覚えておけばほとんど場合の入力に対応できます。

〇文字列を受け取る
例
abc

s=input()

〇ひとつの整数を受け取る
例
3

n=int(input())

〇横に並んだ複数の文字列を受け取る
例
abc arc agc

s,t,u=input().split()

li=input().split()

前者は３つの変数s、t、uを用意してabc、arc、agcを受け取っているのに対し、後者はリストliを用意してまとめて受け取っている。場面によって使い分ける。

〇横に並んだ複数の整数を受け取る
例
1 2 3

a,b,c=map(int,input().split())

li=list(map(int,input().split()))

前者は３つの変数a、b、cを用意して1、2、3を受け取っているのに対し、後者はリストliを用意してまとめて受け取っている。場面によって使い分ける。

〇縦に並んだ複数の文字列を受け取る
例（文字列の数である3が先頭に与えられる）
3
abc
arc
agc

n=int(input())
li=[input() for i in range(n)]

nは文字列の数である3を受け取っている。そしてリストliを用意し、リスト内包表記を使って3つの文字列をまとめて受け取っている。

〇縦に並んだ複数の整数を受け取る
例（整数の数である3が先頭に与えられる）
3
100
200
300

n=int(input())
li=[int(input()) for i in range(n)]

nは文字列の数である3を受け取っている。そしてリストliを用意し、リスト内包表記を使って3つの整数をまとめて受け取っている。

〇複数列、複数行の整数を受け取る
例（行数である3が先頭に与えられる）
3
100 200
300 400
500 600

n=int(input())
for i in range(n):
    a,b=map(int,input().split())

n=int(input())
li=[list(map(int,input().split())) for i in range(n)]

前者はfor文を回して変数a、bに入力していく方法。後者はリスト内包表記を使って二次元配列を作る方法。場面によって使い分ける。