実装して学ぶ基本データ構造~Stackの応用編~

面積計算

これまでの知識を活用した応用問題を解いていきたいと思います．
同じ問題を解きたい方はAOJの問題から解いてみてください
二次元上に対して，"\", "/", ""の三種類の入力が入ってきます．この三種類の入力によってできた図形に対して，雨が降った時に溜まる面積を求めるというのが本問題の目的になります．また，求めた面積の総和に加えて，図形の中にできた水溜りの数とその面積を求める必要があります．
対応する問題はALDS1_3_D: Areas on the Cross-Section Diagram
この問題における入力例で問題の解説をしていきたいと思います．
入力 : \///_/\/\\//\///\_\/\//\
出力 : 35
　　　5 4 2 1 19 9
この入力を図にしたものが下図になります．
今回の解き方としては

• "\"が来たら
  1. "\"の入力の位置をリストxに追加する.
• "/" が来たら
  1. リストxから"\"の位置情報を取り出す
  2. "/"の位置と取り出した位置情報の差を計算する．
  3. "\"と"/"の位置の間にすでに面積リストに追加されている要素があればその面積と2の計算結果を足し合わせる．
  4. 確定した計算結果を面積のリストに追加する．

解き方だけではイメージがわからないと思うので，それぞれのタイミングで何をしているかとそのイメージをgifにまとめたので確認してみてください.
それでは，実際に実装していきます．

def area_cal(input_fig):
    stack = []
    # area = [[面積計算が始まった位置, 面積],・・・]
    area = []
    total = 0
    for i in range(len(input_fig)):
        fig = input_fig[i]
        if fig == "\\":
            stack.append(i)
        elif fig == "/" and stack:
            j = stack.pop()
            width = i - j
            total += width
            while area and area[-1][0] > j:
                width += area[-1][1]
                area.pop()
            area.append([i,width])
    return area, total
            
input_fig = input()
area, total = area_cal(input_fig)
print(total)
if area:
    print(len(area),*list(zip(*area))[1])
else:
    print(0)

提出してみたい際はこのままの形で提出すれば通ると思います．また，動作を確認したい方はinput()の部分を自分なりの記号列に書き換えて見ると色々試せると思います.
ぜひ解いてみてください.