Solve A ~ D of yuki coder 247 with python

I participated in yuki coder 247 (https://yukicoder.me/contests/262). I could only solve up to C within the time, but after the end I solved up to D, so I will put up to D.

A If the freshness of the fruit is positive, sort it and rank the positive one. If you only have negative ones, you can choose one of the freshest negative ones.

# A

N, K = list(map(int, input().split()))
As = list(map(int, input().split()))
positives = [a for a in As if a >= 0]
negatives = [a for a in As if a < 0]

if positives:
    selected = sorted(positives)
else:
    selected = [sorted(negatives)[-1]]

selected = selected[-min(K, len(selected)):]
print(sum(selected))

B If $ N_0 = 0, N_ {i + 1} = AN_i + B $ and $ N_ {n} = 0 $, find the minimum number of steps $ n $. This is $ N_ {n + 1} = \ sum_ {i = 1} ^ n A ^ i B = \ frac {1-A ^ {n + 1}} {1-A} B $, so $ B It can be seen only when = 0 $ or $ A = -1 $.

# B

A, B = list(map(int, input().split()))

if B == 0:
    print(1)
elif A == -1:
    print(2)
else:
    print(-1)

C Find out if you get stuck in one of the following:

There exists an integer $ n $ that is $ LK \ leq nM \ leq RK $
K=0

1 can be any difference between the integer part of $ RK / M $ and the integer part of $ LK / M $ being 1 or more, or $ R $ or $ L $ being a multiple of $ M $. The reason for making the difference more than 1 is to consider when the difference between $ LK $ and $ RK $ is less than $ M $.

# C

L, R, M, K = list(map(int, input().split()))

if R * K // M - L * K // M >= 1 or R * L % M == 0 or K == 0:
    print('Yes')
else:
    print('No')

D Let $ z_i $ be the number if it is zero at the $ i $ th time, and $ n_i $ if it is not zero. In this case, the following recurrence formula can be created.

\displaystyle{
\begin{aligned}
z_{i+1} &= (P+1)z_i + 2 n_i,  \\
n_{i+1} &= (P-1)z_i + 2(P-1)n_i.
\end{aligned}
}

About the update formula of $ z_i $ --1 item: Total $ P + 1 $ as below -By multiplying $ 0 \ sim P-1 $, $ P $ street -Add $ 0 $ to get $ 1 $ --2 items: 2 ways in total --One way by adding $ P-j $ for each $ j> 0 $ --One way by multiplying by 0 　 About the update formula of $ n_i $ --1 item: Total $ P-1 $ as below --By adding each $ P-j \ (j \ neq 0) $ to $ z_i $, you can follow $ P-1 $. --2 items: Total $ 2 (P-1) $ as below --For each $ j> 0 $, add $ P-k \ (j \ neq k) $ to get $ P-1 $. --By multiplying each $ j> 0 $ by $ P-k \ (k \ neq 0) $, you can get $ P-1 $.

# D

P, K = list(map(int, input().split()))
mod = 10**9 + 7

z_i, n_i = 1, 0
for i in range(K):
    z_i, n_i = (P + 1) * z_i + 2 * n_i, (P - 1) * z_i + 2 * (P - 1) * n_i
    z_i, n_i = z_i % mod, n_i % mod

print(z_i)