Python 謎テクニック



筆者が知っているPythonの謎テクニック系をまとめました。    
ほとんど役に立たないと思います。  

<details><summary>更新履歴</summary>

2024/09/22 更新:
1. 「英小文字をソートして列挙」のコードを訂正しました。
大変失礼いたしました。  

2. 「SortedList速度比較」のBinaryTrieのコードを高速化しました。  

</details>


### テンプレ

筆者は使わないのですが、入力受取の補助関数を定義すると楽です。  
たとえば `li = lambda: list(map(int, input().split()))`とすれば、`li()`の呼び出しでリストの受け取りができます。  



### 要素の列挙

$N$未満の非負整数の列挙は`list( range(N) )`でできます。  
英小文字の列挙はキーボードをがーっとやればよいです。  
```python: くぁｗせｄｒｆｔｇｙふじこｌｐ(訂正後)
small_char = 'qawsedrftgyhujikolpzxcvbnm'
len( {*small_char} )  #26
sorted_small_char = sorted( [*small_char] )  #['a', 'b', ･･･ , 'z']
```

<details><summary>2024/09/22 訂正部分</summary>

3行目のアンパッキング時にエラーが出るコードを掲載しておりました。  
(訂正前) `sorted(*small_char)`
(訂正後) `sorted( [*small_char] )`
大変失礼いたしました。(終)  

```python: くぁｗせｄｒｆｔｇｙふじこｌｐ(訂正前)
small_char = 'qawsedrftgyhujikolpzxcvbnm'
len( {*small_char} )  #26
sorted_small_char = sorted(*small_char)  #['a', 'b', ･･･ , 'z']
```
</details>


### 累積和

> 長さ$N$の数列$A$の累積和配列を作成せよ。
> 入力例: $N = 5, A = [3, 1, 4, 1, 5]$
> 出力例: $[0, 3, 4, 8, 9, 14]$

`sum( A[Lt: Rt] ) = C[Rt] - C[Lt]` となるような配列$C$を作成します。  
基本的には以下の書き方でよいと思います。
```python: 累積和 基本
C = [0] * (N + 1)
for i in range(N):
    C[i + 1] = C[i] + A[i]
    
print(C)  #[0, 3, 4, 8, 9, 14]
```

セイウチ演算子`:=`を用いることで1行で書けます。  
ただし、 __定義した`Ci`はスコープ外に残り続けるので注意してください。__  


```python: 累積和 セイウチ演算子
C = [Ci := 0] + [Ci := Ci + Ai for Ai in A]

print(C)  #[0, 3, 4, 8, 9, 14]
print(Ci)  #14: 内包表記で用いた Ai と異なり、スコープ外に残る
print(Ai)  #NameError: name 'Ai' is not defined.
```

添字の降順に累積和を取らないといけない問題もあります。  
Fenwick Treeで楽をするのが一番ですが、あえて累積和で解く場合は以下のように$A_i$の減算で処理すると都合が良いです。   
__昇順の累積和とは添字が異なります。__ 式中の添字には注意してください。 

```python: 累積和 iの降順
C = [0] * (N + 1)
for i in range(N - 1, -1, -1):
    C[i] = C[i + 1] - A[i]  #変更点: 左辺がC[i], 右辺がC[i + 1]
print(C)  #[-14, -11, -10, -6, -5, 0]
```

### DP系

> 長さ$N$の配列$A$と定数$d$を与える。
> 各$i = 0, 1, ･･･ , N - d$について、$A_i + ･･･ + A_{i + d - 1}$を求めよ。  
> 入力例: $N = 5, A = [3, 1, 4, 1, 5], d = 3$
> 出力例: $[8, 6, 10]$  (3 + 1 + 4, 1 + 4 + 1, 4 + 1 + 5)

添字管理に困りがちですが、`enumerate`を組み合わせると楽に実装できます。  
`range`の値を調整すれば半開区間から閉区間に修正もできます。  

```python: 区間幅dの尺取法
#半開区間[Lt, Rt)
for Lt, Rt in enumerate( range(d, N + 1) ):
    print(Lt, Rt)  #(0, 3), (1, 4), (2, 5)

#閉区間[Lt, Rt]
for Lt, Rt in enumerate( range(d - 1, N) ):
    print(Lt, Rt)  #(0, 2), (1, 3), (2, 4)
```

メモリがきつい問題では2配列を使い回すテクニックを用いることがありますが、初期化の際は毎回代入したほうがメモリに優しいです。  
原理は謎です。ふしぎです。  
ただしそのぶん実行時間が悪くなるのでバランスに注意してください。  
```python: 限界メモリ節約術
inf = 10 ** 18
X, Y = 10, 5
DP1 = [[inf] * Y for _ in range(X)]  #今のDP
DP2 = [[inf] * Y for _ in range(X)]  #次のDP
for _ in range(100):
    #DP2を初期化したい場合
    for x in range(X):
        for y in range(Y):
            DP2[x][y] = inf  #毎回代入するとメモリに優しい

    #毎回定義するとメモリを消費しがち
    del DP2[:], DP2  #delしてもあまり変わらない
    DP2 = [[inf] * Y for _ in range(X)]  #毎回定義し直すとメモリを使う
```

### 座標圧縮

> 数列$A$の重複を削除し、昇順に出力せよ。
> 入力例: $A = [3, 1, 4, 1, 5]$
> 出力例: $[1, 3, 4, 5]$

`sorted( set(A) )`でできますが、setは遅いです。  
ソートしてから隣接2項の一致判定をした方が早いです。  
```python: 座標圧縮
B = sorted(A)
C = [ B[0] ] if len(B) > 0 else []
for Bi in B:
    if C[-1] != Bi: C.append(Bi)

print(C)
```


### 連想配列

> 全要素が$0$で初期化された連想配列$D$を与える。  
> $D_{K_i} += V_i$のクエリを処理し、最終的な$D$を出力せよ。  
> 入力例: $K = [9, 9, 8], V = [2, 4, 4]$
> 出力例: $(K, V) = (9, 6), (8, 4)$

`defaultdict(int)`で解けますが実行時間が不安定で、遅いときは通常のdictの倍は遅いです。  
定数倍を抑えたい場面では、通常のdictを使いつつキーのアクセス回数を極力抑える実装が必要となります。  
代案はいくつかありますが、`try and except`が比較的安定して早い印象です。  
`if Ki in D`系統はその次で、`D.get(Ki, 0)`はちょっと遅めです。  


```python: defaultdict 代案
#1. Ki in Dで毎回判定: アクセス2回
for Ki, Vi in query:
    if Ki in D: D[Ki] += Vi  #Ki in Dで1回、D[Ki]への操作で1回アクセス
    else: D[Ki] = Vi

#2. ワンライナー化: アクセス2回
for Ki, Vi in query:
    D[Ki] = (D[Ki] if Ki in D else 0) + Vi  #D[Ki]は2回あるがアクセスは1回

#3. try and except: アクセス1～2回
for Ki, Vi in query:
    try: D[Ki] += Vi  #1回目のアクセス
    except: D[Ki] = Vi  #キーが存在しなかった場合のみ、追加で1回アクセス

#4. D.get(key, default): アクセス2回
for Ki, Vi in query:
    D[Ki] = D.get(Ki, 0) + Vi  #D[Ki]とD.getで2回アクセス
```

<details><summary>アクセス回数の推定法</summary>

ハッシュテーブルに入力する値のハッシュ値を強制的に固定すれば、全要素が衝突する連想配列を作れます。  
```python: ハッシュ値を固定
class malice_int(int):  #intを継承した子クラスを作成
    def __init__(self, v):
        int.__init__(v)
    def __hash__(self): return -1  #ハッシュ値を-1に固定
def conflict_dict(N = 10000):
    return {malice_int(Ki): 0 for Ki in range(N)}
```

なお`conflict_dict`は生成するだけで$Θ(N^2)$かかるので、$N$の上げすぎには注意してください。  
この辞書にハッシュ値が$-1$のキーでアクセスし、処理にかかった時間からアクセス回数を推定しました。(終)

</details>


__(注記: 高度な内容を取り扱います。AtCoderではほぼ必要ありません)__

<details><summary>敵対的入力</summary>

以下の記事の通り、Pythonの連想配列は敵対的入力に弱いです。  

https://qiita.com/recuraki/items/086e36ad124598627c41
https://codeforces.com/blog/entry/101817

```python: dictの撃墜例
def destroy_dict(logN = 16):
    from time import perf_counter
    if logN % 2 == 1:  #1. logNを2の倍数に調整
        logN += 1
    N = 1 << max(4, logN)  #2. dictのサイズを決定  4冪にする
    X = [N + 1] + [Ti := 6 % N]  #3. 1を検索したときの連鎖先を先に埋める
    X += [Ti := (Ti * 5 + 1) % N for _ in range(N * 3 // 5 - 2)]   #約60%を埋める
    D = {Xi: 1 for Xi in X}  #4. dictを作成  ここまでおよそ10ms程度
    T = perf_counter()
    for _ in range(10000):  #1の衝突クエリ数
        D[1] = 1
    return perf_counter() - T
```

codeforcesのブログの通り、継承して`__hash__()`関数を直接いじるのもひとつの対策です。  
ただし型変更のコストがかかるので、手動でxorを取るのが確実でしょうか。(終)  

</details>









### BFS・DFS

> 頂点$0$ ～ $N - 1$からなる$N$頂点の根つき木がある。根は頂点$0$である。  
> 隣接リストを配列$G$として与えるので、各頂点の部分木の大きさを求めよ。  
> 入力例: $N = 4, G = [[1], [0, 2], [1, 3], [2]]$
> 出力例: $[4, 3, 2, 1]$  (パスグラフです)

リストでBFSをするとBFSの到達履歴を残すことができます。  
トポロジカルソートをしたい場合に便利かもしれません。  

```python: リストでBFS
size = [1] * N
stack = [(0, -1)]  #(今の頂点, 親の頂点。根の親は存在しない頂点: -1とする)
for now, back in stack:
    for nxt in G[now]:
        if nxt != back:  #移動先が親でなければ、stackに追加
            stack.append((nxt, now))
print(stack)  #[(0, -1), (1, 0), (2, 1), (3, 2)]

while stack:  #トポロジカルソートの逆順で処理
    now, back = stack.pop()
    for nxt in G[now]:
        if nxt != back:
            size[now] += size[nxt]
print(size)
```

スタックをforループではなくpopで回せば非再帰DFSになります。  
詳しくはKiri8128さんの記事をご覧ください。  

https://qiita.com/Kiri8128/items/2b0023bed9af642c751c


最短距離を求めるBFS(01-BFS)の場合はdequeが必要です。  
`from collections import deque`を避けたい人はstack2本を交換しながらBFSすればよいです。  
```python: 01-BFS
dist = [inf := 10 ** 9] * N
dist[0] = 0  #初期化
Q, R = [(0, -1)], []
while Q or R:
    if not Q:  #スタックを交換
        Q, R = R, Q
    while Q:  #Qから取りだす
        now, back = Q.pop()
        for nxt in G[now]:
            if dist[nxt] > dist[now] + 1:
                dist[nxt] = dist[now] + 1
                R.append(nxt)  #Rに追加する
```

Dial's algorithmを使う場合、ダイアルサイズは2冪にすると除算処理が有利です。    
ダイアルの終了条件はwhile文で書いて、同一ダイアル中でpushが起きた回数が1回以上なら継続とするのが丁寧です。  
変数管理が嫌ならダイアル回数の上界を決め打ったforループを書いて祈りましょう。  

```python: Dial algorithm 主要部
inf = (N - 1) * 999  #N - 1辺 * 最大距離999 が最悪ケースと仮定する
dist = [0] + [inf] * (N - 1)
Q = [[] for _ in range(1 << 10)]  #ダイアル
Q[0].append(0)  #頂点0を追加
mask = (1 << 10) - 1  #除算用
for _ in range( - (- inf >> 10)):  #ループ回数を決め打ち
    for d in range(1 << 10):
        while Q[d]:
            now = Q[d].pop()
            if dist[now] & mask != d: continue
```
[ZONe2021E 提出例(Dial algoほか 継続判定はforループ + お祈り)](https://atcoder.jp/contests/zone2021/submissions/57909913)



### ダイクストラ法

`heapq`モジュールよりも一点更新セグメント木のほうが早い場合があります。  
ふしぎです。  

実装は一点更新・区間最小値のセグメント木を建てて、最短距離が確定したらセグメント木のノードをinfに修正すればよいです。  
区間積取得の関数は要らないので簡潔に書けると思います。  

<details><summary>ダイクストラ用のセグメント木 実装例</summary>

ちょっと便利ですが、あえてライブラリ化する必要はないと思います。  

各頂点の距離を内部配列に持つようにしたほか、新たに`chmin`関数を定義しました。  
距離の更新を`chmin`で行うことで、更新の計算量が$Θ(logN)$から$O(logN)$になります。


```python: ダイクストラ用セグメント木
class dijkstra_SegTree:  #ノードに値を直接代入する版
    def __init__(self, N, inf):  #inf: minの単位元  10 ^ 18など
        self._N = N
        self._logN = N.bit_length()
        self._size = 1 << self._logN
        self._dist = [inf] * N
        self._node = [inf for _ in range(self._size << 1)]
    def build(self, A):  #distをarray Aで初期化  O(N)
        assert len(A) == self._N
        for i, v in enumerate(A):
            self._dist[i] = self._node[i | self._size] = v
        for i in range(self._size - 1, 0, -1):
            self._node[i] = min(self._node[i << 1], self._node[i << 1 | 1])
    def chmin(self, i, v):  #dist[i] ← min(dist[i], v)  O(logN)
        if v < self._dist[i]:
            self._dist[i] = self._node[now := i | self._size] = v
            for _ in range(self._logN):
                if v < self._node[now := now >> 1]:
                    self._node[now] = v
                else:
                    return
    def get(self):  #【まだpopしていない】距離最小の頂点を取り出す  なければ-1  Θ(logN)
        if (v := self._node[now := 1]) == self._node[0]:
            return -1
        for _ in range(self._logN):
            now = e if self._node[e := now << 1] == v else e | 1
        m = self._node[i := now] = self._node[0]
        for _ in range(self._logN):
            if self._node[now ^ 1] < m:
                m = self._node[now ^ 1]
            self._node[now := now >> 1] = m
        return i ^ self._size
    def dist(self, i):
        return self._dist[i]
```

Fenwick Treeでchminを取る版も作りましたが遅かったです。(終)

</details>

[ABC362D 提出例(一点更新セグメント木)](https://atcoder.jp/contests/abc362/submissions/57904420)
[ABC364G 提出例(ダイクストラ用に専用関数を追加、PyPy3 fastest)](https://atcoder.jp/contests/abc364/submissions/57902640)


### ダブリング

> $B^E$を$M$で割った値を求めよ。
> 入力例: $B = 3, E = 14, M = 15$
> 出力例: $9$  ($3^{14} = 4782969 = 318864 × 15 + 9$ )

`bin`関数で上の桁からダブリングするか、ビットシフトで下の桁から2冪を積むかすればよいと思います。  
ただしっくりこないので、もっと楽な実装が知りたいです。  

```python: ダブリング
mul = lambda x, y, M: x * y % M
def doubling_1(B, E, M):  #二進数表記で上の桁から計算
    A = 1  #単位元
    for b in bin(E)[2:]:
        A = mul(A, A, M)
        if b == '1': A = mul(A, B, M)
    return A

def doubling_2(B, E, M):  #ビットシフトで下の桁から計算
    A, C = 1, B  #単位元と2冪作用させる種
    while E > 0:
        if E & 1 == 1: A = mul(A, C, M)
        E >>= 1
        C = mul(C, C, M)
    return A
```

最小共通祖先やSparse Tableで空間$O(NlogN)$のテーブルを作るタイプのダブリングは、配列を`DP[ logN ][N]`の順で定義しましょう。  
`DP[N][ logN ]`の順だと定数倍が悪化します。体感1.3倍程度になります。  



### SortedList関連

tatyamさんのSortedSetの1強です。  
オフライン(クエリ先読みができる問題)ならFenwick Treeと二分探索を使ってもよいです。  

`sortedcontainers`の`SortedSet`は内部的に`SortedList`と`set()`を用いていますが、`set`のせいで速度不利を背負いやすいです。  
少し手間は増えますが、`SortedList`だけで解くことを強くおすすめします。  

<details><summary>SortedList系の解説・速度比較</summary>

[ABC329Dの解法メモ](https://qiita.com/navel_tos/items/6aa57275d137dae786f7)で行った速度比較の追記です。  

[ARC155B Abs Abs Function](https://atcoder.jp/contests/arc155/tasks/arc155_b)で実行速度を比較しました。  
集合に入れる要素はすべて整数で、クエリは`add`(要素の追加)と`prev_value, next_value`(ひとつ前/次の要素の取得)です。  
測定回数は1回です。  
[提出例(全プログラム)](https://atcoder.jp/contests/arc155/submissions/58009836)  

<details><summary>2024/09/22 変更部分</summary>

BinaryTrieのコードを変更し、各ノードの最大値・最小値を陽に持たないようにしました。  
これに伴い、メモリ消費と定数倍が改善しています。(終)

</details>

#### オンラインの解法

##### sortedcontainers: SortedSet 1059ms, SortedList 1022ms

AtCoderの言語アップデートで使えるようになったライブラリです。  
計算量はクエリあたり$O(logN)$のはずです。  
使い方はShirotsumeさんのsortedcontainers解説がおすすめです。  
https://qiita.com/Shirotsume/items/706742162db68c481c3c


##### tatyamさんのSortedSet: SortedSet 484ms, SortedMultiset 476ms

tatyamさんの平方分割ベースのライブラリです。  
__非常に早いです。__  
使い方はQiitaに解説があるので、こちらをご覧ください。  
https://qiita.com/tatyam/items/492c70ac4c955c055602


##### 自作のSkipList: SortedSet 1813ms, SortedList 1723ms

MatsuTakuさんの解説を参考に作成しました。  
詳細は以下のリンクをご覧ください。  
https://qiita.com/MatsuTaku/items/12849a0ae1b1d6427d5a

SkipListとは高さを乱択した双方向リストですが、SkipListの挿入規則を値の昇順にすることでSortedSet(SortedList)となります。  
空間計算量は$O(NlogN)$、時間計算量は各操作 期待$O(logN)$, 最悪$O(N)$です。  

定数倍が悪めなので機能制限をかけたほうがよいでしょう。  
特に多重集合化したり、`__getitem__`(小さい方から$i$番目の要素を取得)を期待$O(logN)$で行おうとすると計算量が一気に悪化します。  


##### 動的Fenwick Tree: 1538ms

__要素が(非負)整数である場合のみ使えます。__
動的Fenwick Treeとは連想配列を用いて空間計算量を$O(QlogN)$に抑えたFenwick Treeです。  
`next_value(v)`は$0$以上$v$以下の個数を数えた後、Fenwick Tree内の二分探索で次の要素を探せばよいです。  

負値の入力が予想される場合、Fenwick Treeを2本立てるか適当なoffsetを加算して非負整数で統一した処理を行えばよいです。  
hashmapは自作のものより標準の`dict`を使う方が早かったです。


##### n分木: Binary Trie 987ms, 4分木 1059ms, 8分木 1165ms, 16分木 1515ms

__要素が(非負)整数である場合のみ使えます。__
入力された整数をバケットサイズごとに区切り、文字とみなしてTrie木に挿入します。  
バケットサイズを大きくするとクエリあたりの時間計算量が減りますが、空間計算量と`__getitem__`(小さい方から$i$番目の要素の取得)の計算量が悪化します。  
Fenwick Treeの要領にすれば$O(maxAi)$で抑えられると思いますが未実装です。実装できたら追記します。   


#### オフラインの解法

クエリ先読みができる状況を仮定します。  

##### 座標圧縮 + セグメント木: 1199ms

`prev_value`をセグメント木の区間最小値で計算する手法です。  
あえてこの方法で実装する必要はないと思います。  

##### 座標圧縮 + Fenwick Tree: 569ms

オフラインで解くならこれが一般的かと思います。  
Fenwick Treeに各添字の個数を登録し、Fenwick Tree内の二分探索で`next_value`を取得します。  
値から座標圧縮後の添字に変換するのには`bisect.bisect_left`が便利です。$O(logN)$ですが`bisect`モジュールは爆速なので問題ありません。  

##### 座標圧縮 + UnionFind: 861ms

__`add`操作がない場合のみ使えます。__  
各添字ごとに自身以下の最大の添字を保持しておきます。  
`remove`操作のたびに添字をマージすれば`prev_value, next_value`が処理できます。  
本問は`add`操作だけですが、クエリを逆から解くことで`discard`操作のみと読み替えられるので本手法が適用できます。  


その他のテクニックが見つかれば追記します。(終) 

</details>




### おわりに

おわりです。
$O(NlogN)$系の比較もしたいなとおもいました。  
おもいついたら書きます。