2020拼多多：种树[Python][DFS]

思路：dfs+剪枝。(见代码)
要注意的是同样的思路python由于递归上限的问题过不了，因此需要手动设置sys.setrecursionlimit()。

import sys # python 默认递归深度不超过1000，做dfs会比C吃亏
sys.setrecursionlimit(10000000)# 手动修改深度，如果TLE那就是算法的问题

# 小多采购了 N 个品种的树，每个品种的数量是 Ai (树的总数量恰好为 M)
N = int(input())
arr = list(map(int, input().split())) # sum(arr) = M
M, res = sum(arr), []


def dfs(a, m, pre_tree_class): # m是剩下的坑位；pre_tree_class前一种树的种类
    global res
    if m == 0:
        return True
    for i in range(len(a)):
        tree_class = i+1 # 当前的树的种类
        if a[i] * 2 > m+1: # 剪枝(某一类大于剩下的坑位一半时就不符合题意了)
            return False
        if a[i] > 0 and pre_tree_class != tree_class:
            a[i] -= 1 # 种树：推进一个状态
            res.append(tree_class)
            if dfs(a, m-1, tree_class):
                return True
            res.pop(-1) # 如果没找到解则回退状态
            a[i] += 1
    return False

if dfs(arr, M, 0):
    print(" ".join(map(str, res)))
else:
    print("-")