きなこの精進日記[python]

Timber [Facebook Hacer Cup Qualication Round, Problem C]

問題:
www.facebook.com

概要

数直線上にN 本の木、場所Pi,高さはHi の木が立っている。

数直線上にそって任意の木をそれぞれ左右どちらかに倒すことができる。

 操作後にに端点が同じ位置である木同士は1つの木とまとめることができる。

 伐採後の木の長さで最大の長さを求めよ。

制約

 N <= 4*10**6(テストケース全体で)
 -5*10**8 <= pi <= 5*10**8
1 <=hi <= 5*10**8
同じ位置に複数の木が存在することはない。

考察

 二分探索でやろうとして、p+h ,p-hでソートした時のidx を持ったりしてバグらせたので、解説の方針で実装します。

 まず、答えの列に使わない木は倒していないことにすればいいので、全ての木は左右どちらかに倒すと考えて良い。

 同じ位置に複数の木が存在することはないという条件、木が端点以外で重複して存在してはいけないという条件から、向きは左向きのみ、右向きのみ、右向きの列の右端に左向きの列の左端が被っている場合の3通りに絞られる。

 ここで左向き、右向きのみの木の列の長さを端点の場所をkeyとして座圧して持つdp として計算すれば、最後に端点が一致するもの全てに対して、一致するi,jに対して、ans = max(ans,dp_l[i]+dp_r[j])でもとまる。

実装

import bisect
import sys

read = sys.stdin.buffer.read
input = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines

# step1:座圧して(ソートした順番で)DPを左端と右端両方から行う
# dp_l[i] = max(dp_l[j] + h[i],h[i]) p+h をソートして持っておけばできる。
# dp_r[i] = max(dp_r[j] + h[i],h[i]) p-hをソートして持つ
# step2:全てのi に対して、p[i]+h[i] がp[j]-h[j]となるj があるかを二分探索で探す
# dict で書き換える
# step3:step2で一致するi,jに対して、ans = max(ans,dp_l[i]+dp_r[j])
t = int(input())
for case in range(t):
    n = int(input())
    p = [0] * n
    h = [0] * n
    for i in range(n):
        p[i], h[i] = map(int, input().split())
    p, h = zip(*sorted(zip(p, h)))
    # p で昇順にsort 右に倒すものは昇順に、左に倒すものは降順に見る

    # 左から順番にみていく 左から見たものに
    # dp はkey がpos,value が長さ で置く
    dp_r_dict = {}  #  右に倒れる木の列のなかでpos が末端(右端)列の中で最長の長さ
    dp_l_dict = {}
    for i in range(n):
        # dict[p[i]]があればそれを返し、なければ0を返す
        # もしp[i]で終わるものがあればそれをつなげて伸ばすことができるため
        tmp = dp_l_dict.get(p[i], 0) + h[i]
        if dp_l_dict.get(p[i] + h[i], 0) < tmp:  # 同様
            dp_l_dict[p[i] + h[i]] = tmp
    r_max = 0
    for i in range(n - 1, -1, -1):
        tmp = dp_r_dict.get(p[i], 0) + h[i]
        if dp_r_dict.get(p[i] - h[i], 0) < tmp:
            dp_r_dict[p[i] - h[i]] = tmp
            r_max = max(r_max, dp_r_dict[p[i] - h[i]])
    # step3:dp_l_dict に対して、全てのkey でdp_r_dictのkey と一致するのがあるかを調べmaxをとる
    ans = 0
    for pos, l_len in dp_l_dict.items():
        tmp = l_len + dp_r_dict.get(pos, 0)  # 全て右向きでもOK
        ans = max(ans, tmp)
    # 全て左向きを考慮する
    ans = max(ans, r_max)
    case_str = "Case #" + str(case + 1) + ": " + str(ans)
    print(case_str)

頭の良い別解

友人に教えてもらったやり方でdpを1つにする方法を教えてもらったので紹介

基本的なdp の考え方は同じで、ソートした順番でDPを端から行う。

 ただ、左右の切り方を区別せずdp配列を1つだけもつ。
  それぞれの木を見る際にdpの更新は右に倒す方から更新する(更新を先に右に倒す方から行えば重複することはない)

import bisect
import sys

read = sys.stdin.buffer.read
input = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines

# step1:座圧して(ソートした順番で)DPを左端からのみ行う
# step2: それぞれの木を見る際にdpの更新は右に倒す方から更新する(更新を先に右に倒す方から行えば重複することはない)
t = int(input())
for case in range(t):
    n = int(input())
    p = [0] * n
    h = [0] * n
    for i in range(n):
        p[i], h[i] = map(int, input().split())
    p, h = zip(*sorted(zip(p, h)))
    # dp はkey がpos,value が長さ で置く
    dp_dict = {}  #  左から更新した時に現在場所=key で終わるありうる木の列の中で最大なもの
    ans = 0
    for i in range(n):
        # 木i を右向きに倒す
        tmp = dp_dict.get(p[i], 0) + h[i]
        if dp_dict.get(p[i] + h[i], 0) < tmp:  # 同様
            dp_dict[p[i] + h[i]] = tmp
            ans = max(ans, dp_dict[p[i] + h[i]])
        # 左向きに倒す
        tmp = dp_dict.get(p[i] - h[i], 0) + h[i]
        if dp_dict.get(p[i], 0) < tmp:
            dp_dict[p[i]] = tmp
            ans = max(ans, dp_dict[p[i]])
    case_str = "Case #" + str(case + 1) + ": " + str(ans)
    print(case_str)

手元で実行してみただけだと早くはならなかったが(多分dictがO(1)で重くないため)、コードが簡潔になった。