笔试编程题汇总(12)

"""
get the balls with scores, (N, max step length), ball pos list, score list
input:
4 1
-1 1 -3 3
1 1 3 4
output:
(5, [0, 0, 1, 2, 3])
"""

import sys


def move_one(cur_pos, move, max_pos=20):
    pos = cur_pos + move
    pos = max(pos, -max_pos)
    pos = min(pos, max_pos)
    return pos


def get_status_score_flag(ball_pos, ball_score, D, n, pos, max_pos=20):
    if n == -1:
        score = 0
        if pos == 0:
            flag = True
        else:
            flag = False
        if flag:
            print("status for (-1, {}): {}, {}".format(pos, score, flag))
        return score, flag, [pos]
    # status tranfer
    last_pos_list = [move_one(pos, move_step, max_pos) for move_step in range(-D, D+1)]
    
    move_step_list = [pos-last_pos for last_pos in last_pos_list]
    cur_status_list = []
    track_list = []
    for last_pos, move_step in zip(last_pos_list, move_step_list):
        last_score, last_flag, track = get_status_score_flag(ball_pos, ball_score, D, n-1, last_pos, max_pos)
        if not last_flag:
            continue
        last_status = (last_score, last_flag, n-1, last_pos)
        cur_status = forward_one_step(ball_pos, ball_score, last_status, move_step)
        cur_status_list.append(cur_status)
        track.append(cur_status[-2])
        track_list.append(track)
    if len(cur_status_list)>0:
        max_idx = max(list(range(len(cur_status_list))), key=lambda x:cur_status_list[x][0])
        max_cur_status = cur_status_list[max_idx]
        cur_score, cur_n, cur_pos, cur_flag = cur_status_list[max_idx]
        cur_track = track_list[max_idx]
    else:
        cur_score, cur_flag, cur_track = 0, False, [pos]
    if cur_flag:
        print("status for ({}, {}): {}, {}, {}".format(n, pos, cur_score, cur_flag, cur_track))
    return cur_score, cur_flag, cur_track


def forward_one_step(ball_pos, ball_score, last_status, step):
    print("move: ", step)
    score, flag, n, pos = last_status
    # forward
    pos = pos+step
    n = n+1
    if pos == ball_pos[n]:
        score = score + ball_score[n]
    return (score, n, pos, flag)
    

def get_max_score(ball_pos, ball_score, D, max_pos=20):
    score_list = []
    track_list = []
    for idx in range(-max_pos, max_pos):
        score, flag, track = get_status_score_flag(ball_pos, ball_score, D, len(ball_pos)-1, idx, max_pos)
        if flag:
            score_list.append(score)
            track_list.append(track)
        
    max_idx = max(list(range(len(score_list))), key=lambda x: score_list[x])
    max_score = score_list[max_idx]
    track = track_list[max_idx]
    return max_score, track


if __name__=="__main__":
    N, D = map(int, input().split())

    ball_pos = list(map(int, input().split()))
    ball_score = list(map(int, input().split()))
    print(get_max_score(ball_pos, ball_score, D, max_pos=20))