CodeRun/CodeRun_Boost/7266_stairs/solution.py at main · Melodiz/CodeRun · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
def solution(n, a):
    result = [0] * (n + 1)
    if n == 1:
        result[0] = 1
        result[1] = a[0]
        return result

    dp_plus = [0] * n
    dp_minus = [0] * n
    dp_plus[0] = a[0]
    dp_minus[0] = -a[0]

    for i in range(1, n):
        current_plus = a[i]
        current_minus = -a[i]
        plus_ok = False
        minus_ok = False

        if dp_plus[i-1] <= current_plus:
            plus_ok = True
        if dp_minus[i-1] <= current_plus:
            plus_ok = True

        if dp_plus[i-1] <= current_minus:
            minus_ok = True
        if dp_minus[i-1] <= current_minus:
            minus_ok = True

        if not plus_ok and not minus_ok:
            return result

        if plus_ok and minus_ok:
            dp_plus[i] = current_plus
            dp_minus[i] = current_minus
        elif plus_ok:
            dp_plus[i] = current_plus
            dp_minus[i] = float('inf')
        else:
            dp_plus[i] = float('inf')
            dp_minus[i] = current_minus

    result[0] = 1
    current = float('inf')
    for i in range(n-1, -1, -1):
        if i == n-1:
            if dp_plus[i] != float('inf'):
                result[i+1] = dp_plus[i]
                current = dp_plus[i]
            else:
                result[i+1] = dp_minus[i]
                current = dp_minus[i]
        else:
            if dp_plus[i] <= current:
                result[i+1] = dp_plus[i]
                current = dp_plus[i]
            else:
                result[i+1] = dp_minus[i]
                current = dp_minus[i]

    return result