【算法训练营】刷油漆，n皇后（python实现）

刷油漆

---

描述

有n辆车排成一排，还有m种不同颜色的油漆，其中第i种油漆够涂ai辆车，同时所有油漆恰好能涂完n辆车。若任意两辆相邻的车颜色不能相同，有多少种涂油漆的方案？

输入

第一行包含一个正整数m。

接下来一行包含m个正整数，第i个正整数表示ai。

输出

输出一个整数，表示答案除以23333的余数。

样例1输入

3
2 1 3

样例1输出

10

样例1解释

10个方案分别是：

1 3 1 3 2 3
1 3 2 3 1 3
2 3 1 3 1 3
3 1 2 3 1 3
3 1 3 1 2 3
3 1 3 1 3 2
3 1 3 2 1 3
3 1 3 2 3 1
3 2 1 3 1 3
3 2 3 1 3 1

样例2

请查看下发文件内的sample2_input.txt和sample2_output.txt。

限制

n为ai之和。

对于50%的数据，n ≤ 10；

对于100%的数据，m ≤ 20，ai ≤ 5。

时间：10 sec

空间：512 MB

提示

[注意到ai≤5，所以我们可以将“还能涂1辆车的油漆种类数”、“还能涂2辆车的油漆种类数”、...、“还能涂5辆车的油漆种类数”设计成状态，思考一下便能得到转移。]

代码实现 

N = 21
mo = 23333

# f: records previously calculated answers to reduce redundant calculations
f = [[[[[[-1 for _ in range(6)] for _ in range(N)] for _ in range(N)] for _ in range(N)] for _ in range(N)] for _ in
     range(N)]


# Dynamic programming (memoization strategy search) to find the number of painting schemes
def dp(a, b, c, d, e, last):
    if (a | b | c | d | e) == 0:
        return 1

    if f[a][b][c][d][e][last] != -1:
        return f[a][b][c][d][e][last]

    ret = 0

    if a:
        ret += dp(a - 1, b, c, d, e, 1) * (a - (last == 2))
    if b:
        ret += dp(a + 1, b - 1, c, d, e, 2) * (b - (last == 3))
    if c:
        ret += dp(a, b + 1, c - 1, d, e, 3) * (c - (last == 4))
    if d:
        ret += dp(a, b, c + 1, d - 1, e, 4) * (d - (last == 5))
    if e:
        ret += dp(a, b, c, d + 1, e - 1, 5) * e

    f[a][b][c][d][e][last] = ret % mo
    return f[a][b][c][d][e][last]


# b: b[i] represents how many paint types are enough for i cars
b = [0, 0, 0, 0, 0, 0]


# n cars, m paint types, the ith paint type is enough for ai cars, and all paint types can exactly paint n cars.
# If any two adjacent cars cannot have the same color, how many painting schemes are there?
def getAnswer(m, a):
    global f, b
    f = [[[[[[-1 for _ in range(6)] for _ in range(N)] for _ in range(N)] for _ in range(N)] for _ in range(N)] for _ in
         range(N)]

    for i in range(m):
        b[a[i]] += 1

    return dp(b[1], b[2], b[3], b[4], b[5], 0)


if __name__ == "__main__":
    m = int(input())
    a = list(map(int, input().split()))
    print(getAnswer(m, a))

n皇后

---

描述

n皇后问题：一个n×n的棋盘，在棋盘上摆n个皇后，满足任意两个皇后不能在同一行、同一列或同一斜线上的方案有多少种？

输入

第一行包含一个整数n。

输出

输出一个整数，表示方案数。

样例1输入

4

样例1输出

2

样例2

请查看下发文件内的sample2_input.txt和sample2_output.txt。

限制

一共10个测试点， 第i个测试点的n=i+4。

时间：2 sec

空间：512 MB

提示

python同学注意，标程后两个测试点10s都过不去，故自行打表。

[考察剪枝水平，剪枝剪得好（二进制剪枝）的才能过第10个测试点。]

代码实现 

def test(row, ld, rd, ans, all_one):
    if row != all_one:
        pos = all_one & ~(row | ld | rd)
        while pos:
            p = pos & -pos
            pos -= p
            test(row | p, (ld | p) << 1, (rd | p) >> 1, ans, all_one)
    else:
        ans[0] += 1

def get_answer(n):
    ans = [0]
    all_one = (1 << n) - 1
    test(0, 0, 0, ans, all_one)
    return ans[0]

if __name__ == "__main__":
    n = int(input())
    print(get_answer(n))