[代码]BZOJ 2831 小强的金字塔

Abstract

BZOJ 2831 小强的金字塔

数学题

Source

http://www.lydsy.com/JudgeOnline/problem.php?id=2831

Solution

http://wjmzbmr.com/archives/xiaoqiang_pyramid/

一开始sb，写了三个函数，复杂度直接爆……后来问clj被告知他只写了一个函数……

写成一个函数或者搞三个然后记忆化都可以过。

我果然是sb……

Code

/**************************************************************
    Problem: 2831
    User: ToumaKazusa
    Language: Python
    Result: Accepted
    Time:84 ms
    Memory:6616 kb
****************************************************************/
 
sum = lambda n:n*(n+1)/2
sqrsum=lambda n:n*(n+1)*(2*n+1)/6
 
#[x*((ax+b)/c), (ax+b)/c, ((ax+b)/c)^2]
def f(a, b, c, r):
    if r < 0:
        return [0]*3
 
    if a >= c:
        t = f(a%c, b, c, r)
        a /= c
        return [t[0]+a*sqrsum(r), t[1]+a*sum(r), t[2]+a*a*sqrsum(r)+2*a*t[0]]
    elif b >= c:
        t = f(a, b%c, c, r)
        b /= c
        return [t[0]+b*sum(r), t[1]+b*(r+1), t[2]+b*b*(r+1)+2*b*t[1]]
    else:
        if a == 0:
            return [0]*3
        y = (a*r+b)/c
        t = f(c, c-b-1, a, y-1)
        res = [y*sum(r)-(t[1]+t[2])/2, y*r-t[1], y*y*r-2*t[0]-t[1]]
        return res
 
a, c, b, l, r = map(int, raw_input().split())
print f(a, b, c, r)[0]-f(a, b, c, l-1)[0]