cindy_game.py

cindy_game.py

import functools
import sympy
from scipy import optimize

n = 3 # n 个人
m = 3 # m 个坑
def sol(n,m):

    v = sympy.symbols(" ".join([str(i) for i in range(m)]))

    s = functools.reduce(sympy.Symbol.__add__, v) ** (n-1)

    def delete_terms(s,v):
        return functools.reduce(
            sympy.Symbol.__add__,
            [i[1]*v**i[0][0] for i in sympy.Poly(s, v).all_terms() if i[0][0] is not 1],
        )

    def gen(i):
        return functools.reduce(
            delete_terms,
            [sympy.Poly(s, v[i]).all_terms()[-1][1]]+list(v[i+1:])
        )
    """
    def gen(num):
        return functools.reduce(
            lambda s, var: functools.reduce(
                sympy.Symbol.__add__,
                [i[1]*var**i[0][0] for i in sympy.Poly(s, var).all_terms() if i[0][0] is not 1],
            ),
            [sympy.Poly(s, v[num]).all_terms()[-1][1]]+list(v[num+1:])
        )
    """

    def wrap(i):
        return sympy.Poly(gen(i), *v).as_expr

    print("开始构造")
    wraped = [wrap(i) for i in range(m)]
    print("构造完成")

    def func(x):
        return [
            wraped[i](*x) - wraped[i+1](*x) for i in range(m-1)
        ] + [functools.reduce(float.__add__,x) - 1]

    print("开始求根")
    sol = optimize.root(func, [1./m for _ in range(m)])
    print("求根完成")
    return sol.x