AtCoder Grand Contest 002 DのTLEが取れない

前提・実現したいこと

Python3 で AtCoder Grand Contest 002 の D問題を解いています。
想定解は並列二分探索のようなのですが、せっかくなので部分永続UnionFindでコードを書いてみました。しかし、実行時間制限超過が取れません。ネットに載っている定数倍高速化は大抵試しPyPy 3で提出したのですが、最小で69msの超過まできてもそれ以上の高速化ができません。正答が保証されていない言語で想定解ではないので仕方ないのですが、諦めきれず、この機会に定数倍高速化の知見を得られたら嬉しいと思い質問させていただきます。初質問でご無礼がありましたら申し訳ありません。また、UnionFindを改造した半自作ライブラリなので不適切なところがありましたら指摘していただけると助かります。

発生している問題・エラーメッセージ

[TLE] ; Time Limit Exceeded

該当のソースコード

python
1#!/usr/bin/env python3
2def main():
3    from sys import stdin
4    n, m, *q = map(int, stdin.buffer.read().split())
5    def bisect(s, t):
6        ok, ng = len(s), -1
7        while ok - ng > 1:
8            md = (ok + ng) // 2
9            if s[md] > t:
10                ok = md
11            else:
12                ng = md
13        return ok
14    
15    # PartiallyPersistentUnionFind
16    class PPUnionFind():
17        def __init__(self, n):
18            self.n = n
19            self.parents = [-1] * n
20            self.time = [10 ** 5 + 1] * n
21            r = [None] * n
22            self.number_time = [[0] for _ in r]
23            self.number_dots = [[1] for _ in r]
24
25        def find(self, x, t):
26            while self.time[x] <= t:
27                x = self.parents[x]
28            return x
29            ########### or ###########
30            # if self.time[x] > t:
31            #    return x
32            # return self.find(self.parents[x], t)
33
34        def union(self, x, y, t):
35            x = self.find(x, t)
36            y = self.find(y, t)
37            if x == y:
38                return
39            if self.parents[x] > self.parents[y]:
40                x, y = y, x
41            self.parents[x] += self.parents[y]
42            self.parents[y] = x
43            self.number_time[x] += [t]
44            self.number_dots[x] += [-self.parents[x]]
45            self.time[y] = t
46
47        def size(self, x, y, t):
48            x = self.find(x, t)
49            y = self.find(y, t)
50            a = self.number_dots[x][bisect(self.number_time[x], t) - 1]
51            if x != y:
52                a += self.number_dots[y][bisect(self.number_time[y], t) - 1]
53            return a
54
55    UF = PPUnionFind(n + 1)
56    uni = UF.union
57    for t, (a, b) in enumerate(zip(*[iter(q[:m*2])]*2)):
58        uni(a, b, t + 1)
59
60    lst = []
61    p = lst.append
62    size = UF.size
63    for x, y, z in zip(*[iter(q[m*2 + 1:])]*3):
64        ok, ng = m, 0
65        while ok - ng > 1:
66            md = (ok + ng) // 2
67            if size(x, y, md) >= z:
68                ok = md
69            else:
70                ng = md
71        p(ok)
72    print(*lst, sep="\n")
73
74if __name__ == "__main__":
75    main()
76

試したこと

python
1if __name__ == "__main__"

を使ったり、入力をsysで高速化、メソッドを事前代入しておくなどしました。
また、UnionFindのクラスを外し単に関数にしても高速化はできませんでした。
AtCoderでは2100ms程度でジャッジは自動停止しますが、最も重いケースで少なくとも2100ms以下に収まることをコードテストで確認しています。また、正答は全て得られていることも確認しています。PyPyは二分探索も自作の方が速いことを確認し関数化しています。

補足情報

環境： AtCoder Judge PyPy3

Python 3.2.5 (b2091e973da6, Oct 19 2014, 18:29:55)
[PyPy 2.4.0 with GCC 4.6.3]

行動規範の内容に同意します

回答1件

ベストアンサー

コードテストで試した範囲だと、
二分探索の

Python
1(ok + ng) // 2

を

Python
1(ok + ng) >> 1

に置き換えると通りそうな感じです。

投稿2020/06/03 04:59

yudedako67

総合スコア2047

Orangeat

2020/06/03 05:06

早速試してみたところ綺麗にACで通りました！！！完全にその改善方法は頭から抜け落ちていました... 本当に素早いご回答ありがとうございます！！ずっと足掻いていたのでとても助かりました！！ >> PyPy3 (2.4.0) 得点 1000 コード長 2205 Byte 結果 [AC] 実行時間 1992 ms メモリ 128732 KB

行動規範の内容に同意します