[Python] ABC004 D

問題

D – マーブル

回答

特に名のあるアルゴリズムは使っていない回答です。

r_num, g_num, b_num = map(int,input().split())

# 40点回答
# def count(x):
#     if (x-1)%2 == 0:
#         c = (x-1)*(x+1)//4
#         return c
#     else:
#         c = x ** 2//4
#         return c

# print(count(r)+count(g)+count(b))

# 原点を中心と仮定して、移動コストを求める関数。
# 移動コストはマーブルの始点からの距離nによって、1, 2, 3, 4…nという等差数列になっているので、その総和になる。
def func_move_cost(left_position, num):
    # 全てのマーブルを原点より左に寄せる時のコスト 
    if left_position + num <= 0:
        return (-left_position - left_position - num + 1) * num/2
    # 全てのマーブルを原点より右に寄せる時のコスト
    elif left_position >= 0:
        return (left_position + left_position + num - 1) * num/2
    # 原点より左にleft_position、右にleft_position+num-1に移動させる時のコスト
    else:
        return ((-left_position) + 1) * (-left_position)/2 + (left_position + num)*(left_position + num -1)/2


# 答えはとりあえず大きい値にしておき、最後に最小値を求める。
answer = float('inf')
move_cost = 0

for g_left_position in range (-500, 500):
    move_cost = func_move_cost(g_left_position, g_num)

    # redの左が g_left_position - r_num になる場合と -100 - r_num/2 になる場合がある。
    # 100足して原点にする。
    if g_left_position - r_num < -100 - r_num/2:
        move_cost += func_move_cost(g_left_position - r_num + 100, r_num)
    else:
        move_cost += func_move_cost(-(r_num)/2, r_num)
    
    # blueの左が g_left_position + g_num になる場合と 100 - b_num/2になる場合がある。
    # 100引いて原点にする。
    if g_left_position + g_num > 100 - b_num/2:
        move_cost += func_move_cost(g_left_position + g_num-100, b_num)
    else:
        move_cost += func_move_cost(-b_num/2, b_num)

    answer = min(answer, move_cost)

print(int(answer))

40点取った時点で諦めました。場合分けがうまくできなかった。

メモ