CGL_6_A: Segment Intersections: Manhattan Geometry 続き2

ひととおり書ききらないと動かす気がしないけど
動かせるほど書くと動く気がしないというね

でも必死こいてなんとか書ききった
Point、Segment、Node、Treeの定義は以前のものを使いまわし
解説する気力がないので載せるだけ

from sys import stdin

# x1とx2の間にある現在有効な縦の線分の数を数える

def crosses(x1, x2, T):
    node = T.root
    leftmost = None

    # x1より右にあるノードのなかで最も左にあるものを探す
    while node is not None:
        if x1 < node.key[1].p1.x:
            if leftmost is None or node.key[1].p1.x < leftmost.key[1].p1.x:
                leftmost = node
            node = node.left
        elif x1 == node.key[1].p1.x:
            leftmost = node
            break
        else:
            node = node.right

    # leftmostより右で、x2より左にある線分を数える
    n_crosses = 0
    node = leftmost
    while node is not None and node.key[1].p2.x <= x2:
        n_crosses += 1
        node = node.successor()

    return n_crosses

def main():
    n = int(stdin.readline())
    segments = []

    for i in range(n):

        # 線分のデータを読み込む
        x1, y1, x2, y2 = [int(x) for x in stdin.readline().split()]

        # x1 < x2、y1 < y2になるようにする
        if x1 == x2:
            if y1 > y2:
                y1, y2 = y2, y1
        else:
            if x1 > x2:
                x1, x2 = x2, x1

        # 線分をリストに保存
        seg = Segment(Point(x1, y1), Point(x2, y2))
        if x1 == x2:
            # 縦の線分を始点で登録
            segments.append((y1, 1, seg))
            # 縦の線分を終点で登録
            segments.append((y2, 3, seg))
        else:
            segments.append((y1, 2, seg))

    # y座標の小さい順にソート
    # y座標が同じなら、縦線の始点、横線、縦線の終点の順に
    segments.sort(key=lambda elem: (elem[0], elem[1]))

    T = Tree()
    n_cross = 0

    # 線分をy座標の小さいものから見ていく
    for y, k, seg in segments:
        if k == 1:
            # 縦の線分の始点なら木に登録
            T.insert(Node((seg.p1.x, seg)))
        elif k == 3:
            # 縦の線分の終点なら木から削除
            T.delete(T.find((seg.p1.x, seg)))
        else:
            # 横の線分なら現在有効な縦線との交点を数える
            n_cross += crosses(seg.p1.x, seg.p2.x, T)

    print(n_cross)

なんとか書ききったら意外とだいたい動いて何か所か直したらACも取れた