给你一棵二叉树，它的根为 root 。请你删除 1 条边，使二叉树分裂成两棵子树，且它们子树和的乘积尽可能大。
由于答案可能会很大，请你将结果对 10^9 + 7 取模后再返回。
示例 1：
输入：root = [1,2,3,4,5,6] 输出：110 解释：删除红色的边，得到 2 棵子树，和分别为 11 和 10 。它们的乘积是 110 （11*10）
示例 2：
输入：root = [1,null,2,3,4,null,null,5,6] 输出：90 解释：移除红色的边，得到 2 棵子树，和分别是 15 和 6 。它们的乘积为 90 （15*6）
示例 3：
输入：root = [2,3,9,10,7,8,6,5,4,11,1] 输出：1025
示例 4：
输入：root = [1,1] 输出：1
提示：
登录50000 个节点，且至少有 2 个节点。[1, 10000] 之间。第一次 DFS：计算整棵树的点权和 $\textit{total}$。
第二次 DFS：计算子树点权和 $s$，那么删除当前节点到其父节点的边后，另一部分的和就是 $\textit{total} - s$，二者乘积为
$$
s\cdot(\textit{total} - s)
$$
用其更新答案的最大值。
由于本题保证点权是非负，我们无需判断当前节点是根节点的特殊情况（无父节点），此时上式为 $0$，不影响答案。
如果有负数点权，就需要跳过当前节点是根节点的情况。
也可以在第一次 DFS 时，把子树和保存到一个列表中，第二次只需遍历这个列表，见 Python3 的第一份代码。
###py
class Solution:
def maxProduct(self, root: Optional[TreeNode]) -> int:
def dfs(node: Optional[TreeNode]) -> int:
if node is None:
return 0
s = node.val + dfs(node.left) + dfs(node.right)
sub_sum.append(s)
return s
sub_sum = []
total = dfs(root)
ans = max(s * (total - s) for s in sub_sum)
return ans % 1_000_000_007
###py
class Solution:
def maxProduct(self, root: Optional[TreeNode]) -> int:
def dfs1(node: Optional[TreeNode]) -> int:
if node is None:
return 0
return node.val + dfs1(node.left) + dfs1(node.right)
def dfs2(node: Optional[TreeNode]) -> int:
if node is None:
return 0
s = node.val + dfs2(node.left) + dfs2(node.right)
nonlocal ans
ans = max(ans, s * (total - s))
return s
total = dfs1(root)
ans = 0
dfs2(root)
return ans % 1_000_000_007
###java
class Solution {
private static final int MOD = 1_000_000_007;
private long ans = 0;
public int maxProduct(TreeNode root) {
int total = dfs1(root);
dfs2(root, total);
return (int) (ans % MOD);
}
private int dfs1(TreeNode node) {
if (node == null) {
return 0;
}
return node.val + dfs1(node.left) + dfs1(node.right);
}
private int dfs2(TreeNode node, int total) {
if (node == null) {
return 0;
}
int s = node.val + dfs2(node.left, total) + dfs2(node.right, total);
ans = Math.max(ans, (long) s * (total - s));
return s;
}
}
###cpp
class Solution {
public:
int maxProduct(TreeNode* root) {
auto dfs1 = [&](this auto&& dfs1, TreeNode* node) -> int {
if (node == nullptr) {
return 0;
}
return node->val + dfs1(node->left) + dfs1(node->right);
};
long long total = dfs1(root);
long long ans = 0;
auto dfs2 = [&](this auto&& dfs2, TreeNode* node) -> int {
if (node == nullptr) {
return 0;
}
int s = node->val + dfs2(node->left) + dfs2(node->right);
ans = max(ans, s * (total - s));
return s;
};
dfs2(root);
return ans % 1'000'000'007;
}
};
###c
#define MAX(a, b) ((b) > (a) ? (b) : (a))
int maxProduct(struct TreeNode* root) {
int dfs1(struct TreeNode* node) {
if (node == NULL) {
return 0;
}
return node->val + dfs1(node->left) + dfs1(node->right);
}
long long total = dfs1(root);
long long ans = 0;
int dfs2(struct TreeNode* node) {
if (node == NULL) {
return 0;
}
int s = node->val + dfs2(node->left) + dfs2(node->right);
ans = MAX(ans, s * (total - s));
return s;
}
dfs2(root);
return ans % 1000000007;
}
###go
func maxProduct(root *TreeNode) (ans int) {
var dfs1 func(*TreeNode) int
dfs1 = func(node *TreeNode) int {
if node == nil {
return 0
}
return node.Val + dfs1(node.Left) + dfs1(node.Right)
}
total := dfs1(root)
var dfs2 func(*TreeNode) int
dfs2 = func(node *TreeNode) int {
if node == nil {
return 0
}
s := node.Val + dfs2(node.Left) + dfs2(node.Right)
ans = max(ans, s*(total-s))
return s
}
dfs2(root)
return ans % 1_000_000_007
}
###js
var maxProduct = function(root) {
function dfs1(node) {
if (node === null) {
return 0;
}
return node.val + dfs1(node.left) + dfs1(node.right);
}
function dfs2(node) {
if (node === null) {
return 0;
}
const s = node.val + dfs2(node.left) + dfs2(node.right);
ans = Math.max(ans, s * (total - s));
return s;
}
const total = dfs1(root);
let ans = 0;
dfs2(root);
return ans % 1_000_000_007;
};
###rust
use std::rc::Rc;
use std::cell::RefCell;
impl Solution {
pub fn max_product(root: Option<Rc<RefCell<TreeNode>>>) -> i32 {
fn dfs1(node: &Option<Rc<RefCell<TreeNode>>>) -> i32 {
if let Some(n) = node {
let n = n.borrow();
n.val + dfs1(&n.left) + dfs1(&n.right)
} else {
0
}
}
fn dfs2(node: &Option<Rc<RefCell<TreeNode>>>, total: i32, ans: &mut i64) -> i32 {
if let Some(n) = node {
let n = n.borrow();
let s = n.val + dfs2(&n.left, total, ans) + dfs2(&n.right, total, ans);
*ans = (*ans).max(s as i64 * (total - s) as i64);
s
} else {
0
}
}
let total = dfs1(&root);
let mut ans = 0;
dfs2(&root, total, &mut ans);
(ans % 1_000_000_007) as _
}
}
见下面树题单的「§2.3 自底向上 DFS」。
欢迎关注 B站@灵茶山艾府
分为两个步骤：
1.求整个二叉树的和
2.遍历所有子树和，求最大值
double ans = Double.MIN_VALUE;
double allSum;
double nodeSum;
public int maxProduct(TreeNode root) {
allSum = sum(root);
dfs(root);
return (int)(ans % (int)(1e9 + 7));
}
public double sum(TreeNode node){
if(node == null) return 0;
return node.val + sum(node.left) + sum(node.right);
}
public double dfs(TreeNode node){
if(node == null) return 0 ;
nodeSum = node.val + dfs(node.left) + dfs(node.right);
ans = Math.max(ans, (allSum - nodeSum) * nodeSum);
return nodeSum;
}
###cpp
/**
* Definition for a binary tree node.
* struct TreeNode {
* int val;
* TreeNode *left;
* TreeNode *right;
* TreeNode(int x) : val(x), left(NULL), right(NULL) {}
* };
*/
class Solution {
public:
vector<long> sums;
static const int mod = 1e9 + 7;
int maxProduct(TreeNode* root) {
postOrder(root);
long res = -1;
for (int i = 0; i < sums.size() - 1; ++i) {
// 取最大值时不能取模，应该用long型存结果
res = max(res, sums[i] * (sums.back() - sums[i]));
}
return (int)(res % mod);
}
long postOrder(TreeNode* root) {
if (root == nullptr) return 0;
long res = root->val + postOrder(root->left) + postOrder(root->right);
sums.push_back(res);
return res;
}
};
给你一个二叉树的根节点 root。设根节点位于二叉树的第 1 层，而根节点的子节点位于第 2 层，依此类推。
返回总和 最大 的那一层的层号 x。如果有多层的总和一样大，返回其中 最小 的层号 x。
示例 1：
输入：root = [1,7,0,7,-8,null,null] 输出：2 解释： 第 1 层各元素之和为 1， 第 2 层各元素之和为 7 + 0 = 7， 第 3 层各元素之和为 7 + -8 = -1， 所以我们返回第 2 层的层号，它的层内元素之和最大。
示例 2：
输入：root = [989,null,10250,98693,-89388,null,null,null,-32127] 输出：2
提示：
[1, 104]范围内-105 <= Node.val <= 105本题需要计算每一层的节点值之和，适合用 BFS。原理讲解【基础算法精讲 13】。
维护两个信息：最大层和 $\textit{maxSum}$，以及对应的层号 $\textit{ans}$。如果当前层的和大于 $\textit{maxSum}$，那么更新 $\textit{maxSum}$ 和 $\textit{ans}$。
class Solution:
def maxLevelSum(self, root: Optional[TreeNode]) -> int:
max_sum = -inf
ans = 0
q = [root]
level = 1
while q:
tmp = q
q = []
s = 0
for node in tmp:
s += node.val
if node.left:
q.append(node.left)
if node.right:
q.append(node.right)
if s > max_sum:
max_sum = s
ans = level
level += 1
return ans
class Solution {
public int maxLevelSum(TreeNode root) {
int max_sum = Integer.MIN_VALUE;
int ans = 0;
List<TreeNode> q = new ArrayList<>();
q.add(root);
for (int level = 1; !q.isEmpty(); level++) {
List<TreeNode> tmp = q;
q = new ArrayList<>();
int s = 0;
for (TreeNode node : tmp) {
s += node.val;
if (node.left != null) {
q.add(node.left);
}
if (node.right != null) {
q.add(node.right);
}
}
if (s > max_sum) {
max_sum = s;
ans = level;
}
}
return ans;
}
}
class Solution {
public:
int maxLevelSum(TreeNode* root) {
int max_sum = INT_MIN;
int ans = 0;
vector<TreeNode*> q = {root};
for (int level = 1; !q.empty(); level++) {
auto tmp = q;
q.clear();
int s = 0;
for (auto node : tmp) {
s += node->val;
if (node->left) {
q.push_back(node->left);
}
if (node->right) {
q.push_back(node->right);
}
}
if (s > max_sum) {
max_sum = s;
ans = level;
}
}
return ans;
}
};
func maxLevelSum(root *TreeNode) (ans int) {
maxSum := math.MinInt
q := []*TreeNode{root}
for level := 1; q != nil; level++ {
tmp := q
q = nil
s := 0
for _, node := range tmp {
s += node.Val
if node.Left != nil {
q = append(q, node.Left)
}
if node.Right != nil {
q = append(q, node.Right)
}
}
if s > maxSum {
maxSum = s
ans = level
}
}
return ans
}
创建一个列表 $\textit{rowSum}$，维护每一层的元素和。
在 DFS 的同时，把节点值加到这一层的元素和中。
最后求出最大层和对应的层号。
class Solution:
def maxLevelSum(self, root: Optional[TreeNode]) -> int:
def dfs(node: Optional[TreeNode], level: int) -> None:
if node is None:
return
if len(row_sum) == level: # 首次访问 level 层
row_sum.append(node.val) # 节点值作为层和的初始值
else:
row_sum[level] += node.val
dfs(node.left, level + 1)
dfs(node.right, level + 1)
row_sum = []
dfs(root, 0)
return row_sum.index(max(row_sum)) + 1 # 层号从 1 开始
class Solution {
private void dfs(TreeNode node, int level, List<Integer> rowSum) {
if (node == null) {
return;
}
if (rowSum.size() == level) { // 首次访问 level 层
rowSum.add(node.val); // 节点值作为层和的初始值
} else {
rowSum.set(level, rowSum.get(level) + node.val);
}
dfs(node.left, level + 1, rowSum);
dfs(node.right, level + 1, rowSum);
}
public int maxLevelSum(TreeNode root) {
List<Integer> rowSum = new ArrayList<>();
dfs(root, 0, rowSum);
int ans = 0;
for (int i = 1; i < rowSum.size(); i++) {
if (rowSum.get(i) > rowSum.get(ans)) {
ans = i;
}
}
return ans + 1; // 层号从 1 开始
}
}
class Solution {
public:
int maxLevelSum(TreeNode* root) {
vector<int> row_sum;
auto dfs = [&](this auto&& dfs, TreeNode* node, int level) -> void {
if (node == nullptr) {
return;
}
if (row_sum.size() == level) { // 首次访问 level 层
row_sum.push_back(node->val); // 节点值作为层和的初始值
} else {
row_sum[level] += node->val;
}
dfs(node->left, level + 1);
dfs(node->right, level + 1);
};
dfs(root, 0);
return ranges::max_element(row_sum) - row_sum.begin() + 1; // 层号从 1 开始
}
};
func maxLevelSum(root *TreeNode) (ans int) {
rowSum := []int{}
var dfs func(*TreeNode, int)
dfs = func(node *TreeNode, level int) {
if node == nil {
return
}
if len(rowSum) == level { // 首次访问 level 层
rowSum = append(rowSum, node.Val) // 节点值作为层和的初始值
} else {
rowSum[level] += node.Val
}
dfs(node.Left, level+1)
dfs(node.Right, level+1)
}
dfs(root, 0)
for i, s := range rowSum {
if s > rowSum[ans] {
ans = i
}
}
return ans + 1 // 层号从 1 开始
}
见下面树题单的「§2.13 二叉树 BFS」。
欢迎关注 B站@灵茶山艾府
根据题意，使用 BFS 进行层序遍历即可。
每次以「层」为单位进行拓展，统计该层的元素和，维护处理过程中的最大值层数和，以及层深度。
代码：
###Java
class Solution {
public int maxLevelSum(TreeNode root) {
Deque<TreeNode> d = new ArrayDeque<>();
int max = -0x3f3f3f3f, depth = 1, ans = 0;
d.addLast(root);
while (!d.isEmpty()) {
int sz = d.size(), cur = 0;
while (sz-- > 0) {
TreeNode t = d.pollFirst();
if (t.left != null) d.addLast(t.left);
if (t.right != null) d.addLast(t.right);
cur += t.val;
}
if (cur > max) {
max = cur; ans = depth;
}
depth++;
}
return ans;
}
}
###TypeScript
function maxLevelSum(root: TreeNode | null): number {
const d: TreeNode[] = new Array<TreeNode>()
let he = 0, ta = 0
d[ta++] = root
let max = -0x3f3f3f3f, depth = 1, ans = 0
while (he < ta) {
let sz = ta - he, cur = 0
while (sz-- > 0) {
const t = d[he++]
if (t.left != null) d[ta++] = t.left
if (t.right != null) d[ta++] = t.right
cur += t.val
}
if (cur > max) {
max = cur; ans = depth
}
depth++
}
return ans
};
如果有帮助到你，请给题解点个赞和收藏，让更多的人看到 ~ ("▔□▔)/
也欢迎你 关注我 和 加入我们的「组队打卡」小群 ，提供写「证明」&「思路」的高质量题解。
所有题解已经加入 刷题指南，欢迎 star 哦 ~
我们可以采用深度优先搜索来遍历这棵二叉树，递归的同时记录当前的层号。
相比哈希表，这里我们采用效率更高的动态数组来维护每一层的元素之和，如果当前层号达到了数组的长度，则将节点元素添加到数组末尾，否则更新对应层号的元素之和。
然后遍历数组，找到元素之和最大，且层号最小的元素。
###Python
class Solution:
def maxLevelSum(self, root: Optional[TreeNode]) -> int:
sum = []
def dfs(node: TreeNode, level: int) -> None:
if level == len(sum):
sum.append(node.val)
else:
sum[level] += node.val
if node.left:
dfs(node.left, level + 1)
if node.right:
dfs(node.right, level + 1)
dfs(root, 0)
return sum.index(max(sum)) + 1 # 层号从 1 开始
###C++
class Solution {
vector<int> sum;
void dfs(TreeNode *node, int level) {
if (level == sum.size()) {
sum.push_back(node->val);
} else {
sum[level] += node->val;
}
if (node->left) {
dfs(node->left, level + 1);
}
if (node->right) {
dfs(node->right, level + 1);
}
}
public:
int maxLevelSum(TreeNode *root) {
dfs(root, 0);
int ans = 0;
for (int i = 0; i < sum.size(); ++i) {
if (sum[i] > sum[ans]) {
ans = i;
}
}
return ans + 1; // 层号从 1 开始
}
};
###Java
class Solution {
private List<Integer> sum = new ArrayList<Integer>();
public int maxLevelSum(TreeNode root) {
dfs(root, 0);
int ans = 0;
for (int i = 0; i < sum.size(); ++i) {
if (sum.get(i) > sum.get(ans)) {
ans = i;
}
}
return ans + 1; // 层号从 1 开始
}
private void dfs(TreeNode node, int level) {
if (level == sum.size()) {
sum.add(node.val);
} else {
sum.set(level, sum.get(level) + node.val);
}
if (node.left != null) {
dfs(node.left, level + 1);
}
if (node.right != null) {
dfs(node.right, level + 1);
}
}
}
###C#
public class Solution {
private IList<int> sum = new List<int>();
public int MaxLevelSum(TreeNode root) {
DFS(root, 0);
int ans = 0;
for (int i = 0; i < sum.Count; ++i) {
if (sum[i] > sum[ans]) {
ans = i;
}
}
return ans + 1; // 层号从 1 开始
}
private void DFS(TreeNode node, int level) {
if (level == sum.Count) {
sum.Add(node.val);
} else {
sum[level] += node.val;
}
if (node.left != null) {
DFS(node.left, level + 1);
}
if (node.right != null) {
DFS(node.right, level + 1);
}
}
}
###go
func maxLevelSum(root *TreeNode) (ans int) {
sum := []int{}
var dfs func(*TreeNode, int)
dfs = func(node *TreeNode, level int) {
if level == len(sum) {
sum = append(sum, node.Val)
} else {
sum[level] += node.Val
}
if node.Left != nil {
dfs(node.Left, level+1)
}
if node.Right != nil {
dfs(node.Right, level+1)
}
}
dfs(root, 0)
for i, s := range sum {
if s > sum[ans] {
ans = i
}
}
return ans + 1 // 层号从 1 开始
}
###C
#define MAX_NODE_SIZE 10000
void dfs(struct TreeNode *node, int level, int *sum, int *sumSize) {
if (level == *sumSize) {
sum[*sumSize] = node->val;
(*sumSize)++;
} else {
sum[level] += node->val;
}
if (node->left) {
dfs(node->left, level + 1, sum, sumSize);
}
if (node->right) {
dfs(node->right, level + 1, sum, sumSize);
}
}
int maxLevelSum(struct TreeNode* root) {
int *sum = (int *)malloc(sizeof(int) * MAX_NODE_SIZE);
int sumSize = 0;
dfs(root, 0, sum, &sumSize);
int ans = 0;
for (int i = 0; i < sumSize; ++i) {
if (sum[i] > sum[ans]) {
ans = i;
}
}
return ans + 1; // 层号从 1 开始
}
###JavaScript
var maxLevelSum = function(root) {
const sum = [];
const dfs = (node, level) => {
if (level === sum.length) {
sum.push(node.val);
} else {
sum.splice(level, 1, sum[level] + node.val);
}
if (node.left) {
dfs(node.left, level + 1);
}
if (node.right) {
dfs(node.right, level + 1);
}
}
dfs(root, 0);
let ans = 0;
for (let i = 0; i < sum.length; ++i) {
if (sum[i] > sum[ans]) {
ans = i;
}
}
return ans + 1; // 层号从 1 开始
};
###TypeScript
function maxLevelSum(root: TreeNode | null): number {
const sum: number[] = [];
const dfs = (node: TreeNode, level: number): void => {
if (level === sum.length) {
sum.push(node.val);
} else {
sum[level] += node.val;
}
if (node.left) {
dfs(node.left, level + 1);
}
if (node.right) {
dfs(node.right, level + 1);
}
};
if (root) {
dfs(root, 0);
}
let maxIndex = 0;
let maxSum = -Infinity;
for (let i = 0; i < sum.length; i++) {
if (sum[i] > maxSum) {
maxSum = sum[i];
maxIndex = i;
}
}
return maxIndex + 1; // 层号从 1 开始
};
###Rust
use std::rc::Rc;
use std::cell::RefCell;
use std::collections::VecDeque;
impl Solution {
pub fn max_level_sum(root: Option<Rc<RefCell<TreeNode>>>) -> i32 {
let mut sum = Vec::new();
// 使用 DFS 递归遍历
fn dfs(node: &Rc<RefCell<TreeNode>>, level: usize, sum: &mut Vec<i32>) {
let node_ref = node.borrow();
if level == sum.len() {
sum.push(node_ref.val);
} else {
sum[level] += node_ref.val;
}
if let Some(left) = &node_ref.left {
dfs(left, level + 1, sum);
}
if let Some(right) = &node_ref.right {
dfs(right, level + 1, sum);
}
}
if let Some(root_node) = root {
dfs(&root_node, 0, &mut sum);
}
// 找到最大和的层级
let mut max_index = 0;
let mut max_sum = i32::MIN;
for (i, &s) in sum.iter().enumerate() {
if s > max_sum {
max_sum = s;
max_index = i;
}
}
(max_index + 1) as i32 // 层号从 1 开始
}
}
复杂度分析
时间复杂度：$O(n)$，其中 $n$ 是二叉树的节点个数。
空间复杂度：$O(n)$。最坏情况下二叉树是一条链，需要 $O(n)$ 的数组空间以及 $O(n)$ 的递归栈空间。
由于计算的是每层的元素之和，用广度优先搜索来遍历这棵树会更加自然。
对于广度优先搜索，我们可以用队列来实现。初始时，队列只包含根节点；然后不断出队，将子节点入队，直到队列为空。
如果直接套用方法一的思路，我们需要在队列中存储节点和节点的层号。另一种做法是一次遍历完一整层的节点，遍历的同时，累加该层的节点的元素之和，同时用这层的节点得到下一层的节点，这种做法不需要记录层号。
为了代码实现的方便，我们可以使用两个动态数组，第一个数组 $q$ 为当前层的节点，第二个数组 $\textit{nq}$ 为下一层的节点。遍历 $q$ 中节点的同时，把子节点加到 $\textit{nq}$ 中。遍历完当前层后，将 $q$ 置为 $\textit{nq}$。
###Python
class Solution:
def maxLevelSum(self, root: Optional[TreeNode]) -> int:
ans, maxSum = 1, root.val
level, q = 1, [root]
while q:
sum, nq = 0, []
for node in q:
sum += node.val
if node.left:
nq.append(node.left)
if node.right:
nq.append(node.right)
if sum > maxSum:
ans, maxSum = level, sum
q = nq
level += 1
return ans
###C++
class Solution {
public:
int maxLevelSum(TreeNode *root) {
int ans = 1, maxSum = root->val;
vector<TreeNode*> q = {root};
for (int level = 1; !q.empty(); ++level) {
vector<TreeNode*> nq;
int sum = 0;
for (auto node : q) {
sum += node->val;
if (node->left) {
nq.emplace_back(node->left);
}
if (node->right) {
nq.emplace_back(node->right);
}
}
if (sum > maxSum) {
maxSum = sum;
ans = level;
}
q = move(nq);
}
return ans;
}
};
###Java
class Solution {
public int maxLevelSum(TreeNode root) {
int ans = 1, maxSum = root.val;
List<TreeNode> q = new ArrayList<TreeNode>();
q.add(root);
for (int level = 1; !q.isEmpty(); ++level) {
List<TreeNode> nq = new ArrayList<TreeNode>();
int sum = 0;
for (TreeNode node : q) {
sum += node.val;
if (node.left != null) {
nq.add(node.left);
}
if (node.right != null) {
nq.add(node.right);
}
}
if (sum > maxSum) {
maxSum = sum;
ans = level;
}
q = nq;
}
return ans;
}
}
###C#
public class Solution {
public int MaxLevelSum(TreeNode root) {
int ans = 1, maxSum = root.val;
IList<TreeNode> q = new List<TreeNode>();
q.Add(root);
for (int level = 1; q.Count > 0; ++level) {
IList<TreeNode> nq = new List<TreeNode>();
int sum = 0;
foreach (TreeNode node in q) {
sum += node.val;
if (node.left != null) {
nq.Add(node.left);
}
if (node.right != null) {
nq.Add(node.right);
}
}
if (sum > maxSum) {
maxSum = sum;
ans = level;
}
q = nq;
}
return ans;
}
}
###go
func maxLevelSum(root *TreeNode) int {
ans, maxSum := 1, root.Val
q := []*TreeNode{root}
for level := 1; len(q) > 0; level++ {
tmp := q
q = nil
sum := 0
for _, node := range tmp {
sum += node.Val
if node.Left != nil {
q = append(q, node.Left)
}
if node.Right != nil {
q = append(q, node.Right)
}
}
if sum > maxSum {
ans, maxSum = level, sum
}
}
return ans
}
###C
#define MAX_NODE_SIZE 10000
int maxLevelSum(struct TreeNode* root){
int ans = 1, maxSum = root->val;
struct TreeNode **q = (struct TreeNode **)malloc(sizeof(struct TreeNode *) * MAX_NODE_SIZE);
struct TreeNode **nq = (struct TreeNode **)malloc(sizeof(struct TreeNode *) * MAX_NODE_SIZE);
int qSize = 0;
q[qSize++] = root;
for (int level = 1; qSize > 0; ++level) {
int sum = 0, nqSize = 0;
for (int i = 0; i < qSize; i++) {
sum += q[i]->val;
if (q[i]->left) {
nq[nqSize++] = q[i]->left;
}
if (q[i]->right) {
nq[nqSize++] = q[i]->right;
}
}
if (sum > maxSum) {
maxSum = sum;
ans = level;
}
struct TreeNode *tmp = q;
q = nq;
nq = tmp;
qSize = nqSize;
}
free(q);
free(nq);
return ans;
}
###JavaScript
var maxLevelSum = function(root) {
let ans = 1, maxSum = root.val;
let q = [];
q.push(root);
for (let level = 1; q.length > 0; ++level) {
const nq = [];
let sum = 0;
for (const node of q) {
sum += node.val;
if (node.left) {
nq.push(node.left);
}
if (node.right) {
nq.push(node.right);
}
}
if (sum > maxSum) {
maxSum = sum;
ans = level;
}
q = nq;
}
return ans;
};
###TypeScript
function maxLevelSum(root: TreeNode | null): number {
if (!root) {
return 1;
}
let ans = 1;
let maxSum = root.val;
let q: TreeNode[] = [root];
for (let level = 1; q.length > 0; level++) {
const nq: TreeNode[] = [];
let sum = 0;
for (const node of q) {
sum += node.val;
if (node.left) {
nq.push(node.left);
}
if (node.right) {
nq.push(node.right);
}
}
if (sum > maxSum) {
maxSum = sum;
ans = level;
}
q = nq;
}
return ans;
}
###Rust
use std::rc::Rc;
use std::cell::RefCell;
use std::collections::VecDeque;
impl Solution {
pub fn max_level_sum(root: Option<Rc<RefCell<TreeNode>>>) -> i32 {
let mut ans = 1;
let mut max_sum = i32::MIN;
let mut level = 1;
let mut queue = VecDeque::new();
if let Some(root_node) = root {
queue.push_back(root_node);
}
while !queue.is_empty() {
let level_size = queue.len();
let mut sum = 0;
for _ in 0..level_size {
let node = queue.pop_front().unwrap();
let node_ref = node.borrow();
sum += node_ref.val;
if let Some(left) = &node_ref.left {
queue.push_back(left.clone());
}
if let Some(right) = &node_ref.right {
queue.push_back(right.clone());
}
}
if sum > max_sum {
max_sum = sum;
ans = level;
}
level += 1;
}
ans
}
}
复杂度分析
时间复杂度：$O(n)$，其中 $n$ 是二叉树的节点个数。
空间复杂度：$O(n)$。最坏情况下，数组中有 $O(n)$ 个节点。
给你一个 n x n 的整数方阵 matrix 。你可以执行以下操作 任意次 ：
matrix 中 相邻 两个元素，并将它们都 乘以 -1 。如果两个元素有 公共边 ，那么它们就是 相邻 的。
你的目的是 最大化 方阵元素的和。请你在执行以上操作之后，返回方阵的 最大 和。
示例 1：
输入：matrix = [[1,-1],[-1,1]] 输出：4 解释：我们可以执行以下操作使和等于 4 ： - 将第一行的 2 个元素乘以 -1 。 - 将第一列的 2 个元素乘以 -1 。
示例 2：
输入：matrix = [[1,2,3],[-1,-2,-3],[1,2,3]] 输出：16 解释：我们可以执行以下操作使和等于 16 ： - 将第二行的最后 2 个元素乘以 -1 。
提示：
n == matrix.length == matrix[i].length2 <= n <= 250-105 <= matrix[i][j] <= 105提示 $1$
为了使得操作后方阵总和最大，我们需要使得负数元素的总和尽可能大。
对于方阵中的两个负数元素，一定存在一系列的操作使得这两个负数元素均变为正数，且其余元素不变。
对于方阵中的一个正数元素和一个负数元素，一定存在一系列的操作使得这两个元素交换正负，且其余元素不变。
提示 $1$ 解释
第一部分是显然的。
对于第二部分，我们可以任意选择一条连接两个负数元素的有向路径，按顺序对路径上（除终点以外）的每个元素和它对应的下一个元素都执行一次操作。最终路径上除了两个端点以外的其他元素都被执行了两次操作，因此数值不变；两个端点元素都被执行了一次操作二变为正数。
由于方阵是网格，因此上述路径一定存在。
对于第三部分，将第二部分中的一个负数更改为正数即可证明。
提示 $2$
如果方阵中存在一个元素为 $0$，另一个元素为负数。那么一定存在一系列的操作使得负数元素变为正数，且其余元素不变。
提示 $2$ 解释
类似 提示 $1$，将一个负数元素更改为 $0$ 即可证明。
提示 $3$
如果方阵中存在 $0$，那么一定可以通过一系列的操作使得方阵中所有元素均为非负数;
如果方阵中不存在 $0$，那么：
如果方阵中有奇数个负数元素，那么一定可以通过一系列的操作使得方阵中只有一个负数元素，且该负数元素可以在任何位置。同时，无论如何操作，方阵中必定存在负数元素。
如果方阵中有偶数个负数元素，那么一定可以通过一系列的操作使得方阵中不存在负数元素。
提示 $3$ 解释
对于第一部分，反复对 $0$ 和负数元素进行 提示 $2$ 的操作即可。
对于第二部分，我们首先可以证明如果方阵不存在 $0$，那么负数元素数量的奇偶性不会改变。然后，我们可以根据 提示 $1$ 构造出一系列操作从而达到对应的要求。
思路与算法
根据 提示 $3$，我们可以按照方阵的元素分为以下几种情况：
方阵中有 $0$，那么最大方阵和即为所有元素的绝对值之和；
方阵中没有 $0$，且负数元素数量为偶数，那么最大方阵和即为所有元素的绝对值之和；
方阵中没有 $0$，且负数元素数量为奇数，那么最大方阵和即为所有元素的绝对值之和减去所有元素最小绝对值的两倍。
其中，第一种情况也可以按照负数元素数量的奇偶性划入后两种情况中（此时最小绝对值一定为 $0$）。
我们遍历方阵，维护负数元素的数量、元素的最小绝对值以及所有元素的绝对值之和。随后，我们按照负数元素数量的奇偶性计算对应的最大元素和并返回。
最后，矩阵所有元素绝对值之和可能超过 $32$ 位整数的上限，因此对于 $\texttt{C++}$ 等语言，需要使用 $64$ 位整数来维护。
代码
###C++
class Solution {
public:
long long maxMatrixSum(vector<vector<int>>& matrix) {
int n = matrix.size();
int cnt = 0; // 负数元素的数量
long long total = 0; // 所有元素的绝对值之和
int mn = INT_MAX; // 方阵元素的最小绝对值
for (int i = 0; i < n; ++i){
for (int j = 0; j < n; ++j){
mn = min(mn, abs(matrix[i][j]));
if (matrix[i][j] < 0){
++cnt;
}
total += abs(matrix[i][j]);
}
}
// 按照负数元素的数量的奇偶性讨论
if (cnt % 2 == 0){
return total;
}
else{
return total - 2 * mn;
}
}
};
###Python
class Solution:
def maxMatrixSum(self, matrix: List[List[int]]) -> int:
n = len(matrix)
cnt = 0 # 负数元素的数量
total = 0 # 所有元素的绝对值之和
mn = float("INF") # 方阵元素的最小绝对值
for i in range(n):
for j in range(n):
mn = min(mn, abs(matrix[i][j]))
if matrix[i][j] < 0:
cnt += 1
total += abs(matrix[i][j])
# 按照负数元素的数量的奇偶性讨论
if cnt % 2 == 0:
return total
else:
return total - 2 * mn
###Java
class Solution {
public long maxMatrixSum(int[][] matrix) {
int n = matrix.length;
int cnt = 0; // 负数元素的数量
long total = 0; // 所有元素的绝对值之和
int mn = Integer.MAX_VALUE; // 方阵元素的最小绝对值
for (int i = 0; i < n; ++i){
for (int j = 0; j < n; ++j){
mn = Math.min(mn, Math.abs(matrix[i][j]));
if (matrix[i][j] < 0){
++cnt;
}
total += Math.abs(matrix[i][j]);
}
}
// 按照负数元素的数量的奇偶性讨论
if (cnt % 2 == 0){
return total;
} else {
return total - 2 * mn;
}
}
}
###C#
public class Solution {
public long MaxMatrixSum(int[][] matrix) {
int n = matrix.Length;
int cnt = 0; // 负数元素的数量
long total = 0; // 所有元素的绝对值之和
int mn = int.MaxValue; // 方阵元素的最小绝对值
for (int i = 0; i < n; ++i){
for (int j = 0; j < n; ++j){
mn = Math.Min(mn, Math.Abs(matrix[i][j]));
if (matrix[i][j] < 0){
++cnt;
}
total += Math.Abs(matrix[i][j]);
}
}
// 按照负数元素的数量的奇偶性讨论
if (cnt % 2 == 0){
return total;
} else{
return total - 2 * mn;
}
}
}
###Go
func maxMatrixSum(matrix [][]int) int64 {
n := len(matrix)
cnt := 0 // 负数元素的数量
total := int64(0) // 所有元素的绝对值之和
mn := 1 << 30 // 方阵元素的最小绝对值
for i := 0; i < n; i++ {
for j := 0; j < n; j++ {
mn = min(mn, abs(matrix[i][j]))
if matrix[i][j] < 0 {
cnt++
}
total += int64(abs(matrix[i][j]))
}
}
// 按照负数元素的数量的奇偶性讨论
if cnt % 2 == 0 {
return total
} else {
return total - int64(2 * mn)
}
}
func abs(x int) int {
if x < 0 {
return -x
}
return x
}
###C
long long maxMatrixSum(int** matrix, int matrixSize, int* matrixColSize) {
int n = matrixSize;
int cnt = 0; // 负数元素的数量
long long total = 0; // 所有元素的绝对值之和
int mn = INT_MAX; // 方阵元素的最小绝对值
for (int i = 0; i < n; ++i) {
for (int j = 0; j < n; ++j) {
int abs_val = abs(matrix[i][j]);
if (abs_val < mn) {
mn = abs_val;
}
if (matrix[i][j] < 0) {
++cnt;
}
total += abs_val;
}
}
// 按照负数元素的数量的奇偶性讨论
if (cnt % 2 == 0) {
return total;
} else {
return total - 2 * mn;
}
}
###JavaScript
var maxMatrixSum = function(matrix) {
const n = matrix.length;
let cnt = 0; // 负数元素的数量
let total = 0; // 所有元素的绝对值之和
let mn = Number.MAX_SAFE_INTEGER; // 方阵元素的最小绝对值
for (let i = 0; i < n; i++) {
for (let j = 0; j < n; j++) {
const absVal = Math.abs(matrix[i][j]);
mn = Math.min(mn, absVal);
if (matrix[i][j] < 0) {
cnt++;
}
total += absVal;
}
}
// 按照负数元素的数量的奇偶性讨论
if (cnt % 2 === 0) {
return total;
} else {
return total - 2 * mn;
}
};
###TypeScript
function maxMatrixSum(matrix: number[][]): number {
const n = matrix.length;
let cnt = 0; // 负数元素的数量
let total = 0; // 所有元素的绝对值之和
let mn = Number.MAX_SAFE_INTEGER; // 方阵元素的最小绝对值
for (let i = 0; i < n; i++) {
for (let j = 0; j < n; j++) {
const absVal = Math.abs(matrix[i][j]);
mn = Math.min(mn, absVal);
if (matrix[i][j] < 0) {
cnt++;
}
total += absVal;
}
}
// 按照负数元素的数量的奇偶性讨论
if (cnt % 2 === 0) {
return total;
} else {
return total - 2 * mn;
}
}
###Rust
impl Solution {
pub fn max_matrix_sum(matrix: Vec<Vec<i32>>) -> i64 {
let n = matrix.len();
let mut cnt = 0; // 负数元素的数量
let mut total: i64 = 0; // 所有元素的绝对值之和
let mut mn = i32::MAX; // 方阵元素的最小绝对值
for i in 0..n {
for j in 0..n {
let abs_val = matrix[i][j].abs();
mn = mn.min(abs_val);
if matrix[i][j] < 0 {
cnt += 1;
}
total += abs_val as i64;
}
}
// 按照负数元素的数量的奇偶性讨论
if cnt % 2 == 0 {
total
} else {
total - 2 * mn as i64
}
}
}
复杂度分析
时间复杂度：$O(mn)$，其中 $m$ 为 $\textit{matrix}$ 的行数，$n$ 为 $\textit{matrix}$ 的列数。
空间复杂度：$O(1)$。
第二题
刚开始我用了深度搜索，结果越写越觉得不对劲，最后发现：
其实如果负数数量是双数，我们总是能把他们都翻成正数
如果负数数量是单数，我们怎么翻都会至少留下一个负数
以上，记录下矩阵里绝对值最小的数即可
（我是笨比）
`
###c++
class Solution {
public:
long long maxMatrixSum(vector<vector<int>>& matrix) {
int minn = abs(matrix[0][0]);
long long ans = 0;
bool flag = false; //记录是负数的个数是单数还是双数
for(auto &vec: matrix)
for(int num: vec){
if(num < 0){ //是负数的话
flag = !flag;
num = -num; //取绝对值
}
ans += num; //记录绝对值的和
minn = min(num, minn); //记录最小值
}
if(flag) //是单数
return ans - 2 * minn;
return ans;//是双数
}
};
虽然题目规定我们只能操作相邻的元素，但我们可以通过多次操作，把任意两个元素都乘以 $-1$。
每次操作，恰好改变两个数的正负号。
多次操作，恰好改变偶数个数的正负号。
分类讨论：
代码实现时，无需特判是否有 $0$。如果有 $0$，那么代码中的 $\textit{mn}=0$，对 $\textit{total}$ 无影响。
class Solution:
def maxMatrixSum(self, matrix: List[List[int]]) -> int:
total = neg_cnt = 0
mn = inf
for row in matrix:
for x in row:
if x < 0:
neg_cnt += 1
x = -x # 先把负数都变成正数
mn = min(mn, x)
total += x
if neg_cnt % 2:
total -= mn * 2 # 给绝对值最小的数添加负号
return total
class Solution {
public long maxMatrixSum(int[][] matrix) {
long total = 0;
int negCnt = 0;
int mn = Integer.MAX_VALUE;
for (int[] row : matrix) {
for (int x : row) {
if (x < 0) {
negCnt++;
x = -x; // 先把负数都变成正数
}
mn = Math.min(mn, x);
total += x;
}
}
if (negCnt % 2 > 0) {
total -= mn * 2; // 给绝对值最小的数添加负号
}
return total;
}
}
class Solution {
public:
long long maxMatrixSum(vector<vector<int>>& matrix) {
long long total = 0;
int neg_cnt = 0;
int mn = INT_MAX;
for (auto& row : matrix) {
for (int x : row) {
if (x < 0) {
neg_cnt++;
x = -x; // 先把负数都变成正数
}
mn = min(mn, x);
total += x;
}
}
if (neg_cnt % 2) {
total -= mn * 2; // 给绝对值最小的数添加负号
}
return total;
}
};
#define MIN(a, b) ((b) < (a) ? (b) : (a))
long long maxMatrixSum(int** matrix, int matrixSize, int* matrixColSize) {
long long total = 0;
int neg_cnt = 0;
int mn = INT_MAX;
for (int i = 0; i < matrixSize; i++) {
for (int j = 0; j < matrixColSize[i]; j++) {
int x = matrix[i][j];
if (x < 0) {
neg_cnt++;
x = -x; // 先把负数都变成正数
}
mn = MIN(mn, x);
total += x;
}
}
if (neg_cnt % 2) {
total -= mn * 2; // 给绝对值最小的数添加负号
}
return total;
}
func maxMatrixSum(matrix [][]int) int64 {
total, negCnt, mn := 0, 0, math.MaxInt
for _, row := range matrix {
for _, x := range row {
if x < 0 {
negCnt++
x = -x // 先把负数都变成正数
}
mn = min(mn, x)
total += x
}
}
if negCnt%2 > 0 {
total -= mn * 2 // 给绝对值最小的数添加负号
}
return int64(total)
}
var maxMatrixSum = function(matrix) {
let total = 0;
let negCnt = 0;
let mn = Infinity;
for (const row of matrix) {
for (let x of row) {
if (x < 0) {
negCnt++;
x = -x; // 先把负数都变成正数
}
mn = Math.min(mn, x);
total += x;
}
}
if (negCnt % 2) {
total -= mn * 2; // 给绝对值最小的数添加负号
}
return total;
};
impl Solution {
pub fn max_matrix_sum(matrix: Vec<Vec<i32>>) -> i64 {
let mut total = 0;
let mut neg_cnt = 0;
let mut mn = i32::MAX;
for row in matrix {
for mut x in row {
if x < 0 {
neg_cnt += 1;
x = -x; // 先把负数都变成正数
}
mn = mn.min(x);
total += x as i64;
}
}
if neg_cnt % 2 > 0 {
total -= (mn * 2) as i64; // 给绝对值最小的数添加负号
}
total
}
}
见下面贪心与思维题单的「§5.2 脑筋急转弯」。
欢迎关注 B站@灵茶山艾府
给你一个整数数组 nums，请你返回该数组中恰有四个因数的这些整数的各因数之和。如果数组中不存在满足题意的整数，则返回 0 。
示例 1：
输入：nums = [21,4,7] 输出：32 解释： 21 有 4 个因数：1, 3, 7, 21 4 有 3 个因数：1, 2, 4 7 有 2 个因数：1, 7 答案仅为 21 的所有因数的和。
示例 2:
输入: nums = [21,21] 输出: 64
示例 3:
输入: nums = [1,2,3,4,5] 输出: 0
提示：
1 <= nums.length <= 1041 <= nums[i] <= 105对于一个数，它有哪些因子？这不能一眼看出来。
但反过来，对于一个数 $i$，$i$ 的倍数都有因子 $i$。
枚举 $i=1,2,3,\ldots,10^5$，以及 $i$ 的倍数 $j$，把 $j$ 的因子个数加一，因子和加上 $i$。
对于 $\textit{nums}$ 中的因子个数为 $4$ 的元素，累加这些元素的因子和，即为答案。
###py
MX = 100_001
divisor_num = [0] * MX
divisor_sum = [0] * MX
for i in range(1, MX):
for j in range(i, MX, i): # 枚举 i 的倍数 j
divisor_num[j] += 1 # i 是 j 的因子
divisor_sum[j] += i
class Solution:
def sumFourDivisors(self, nums: List[int]) -> int:
ans = 0
for x in nums:
if divisor_num[x] == 4:
ans += divisor_sum[x]
return ans
###java
class Solution {
private static final int MX = 100_001;
private static final int[] divisorNum = new int[MX];
private static final int[] divisorSum = new int[MX];
private static boolean initialized = false;
// 这样写比 static block 快
public Solution() {
if (initialized) {
return;
}
initialized = true;
for (int i = 1; i < MX; i++) {
for (int j = i; j < MX; j += i) { // 枚举 i 的倍数 j
divisorNum[j]++; // i 是 j 的因子
divisorSum[j] += i;
}
}
}
public int sumFourDivisors(int[] nums) {
int ans = 0;
for (int x : nums) {
if (divisorNum[x] == 4) {
ans += divisorSum[x];
}
}
return ans;
}
}
###cpp
constexpr int MX = 100'001;
int divisor_num[MX];
int divisor_sum[MX];
int init = [] {
for (int i = 1; i < MX; i++) {
for (int j = i; j < MX; j += i) { // 枚举 i 的倍数 j
divisor_num[j]++; // i 是 j 的因子
divisor_sum[j] += i;
}
}
return 0;
}();
class Solution {
public:
int sumFourDivisors(vector<int>& nums) {
int ans = 0;
for (int x : nums) {
if (divisor_num[x] == 4) {
ans += divisor_sum[x];
}
}
return ans;
}
};
###go
const mx = 100_001
var divisorNum, divisorSum [mx]int
func init() {
for i := 1; i < mx; i++ {
for j := i; j < mx; j += i { // 枚举 i 的倍数 j
divisorNum[j]++ // i 是 j 的因子
divisorSum[j] += i
}
}
}
func sumFourDivisors(nums []int) (ans int) {
for _, x := range nums {
if divisorNum[x] == 4 {
ans += divisorSum[x]
}
}
return
}
设 $U=\max(\textit{nums})$。预处理的时间为 $\mathcal{O}(U\log U)$（调和级数和），空间为 $\mathcal{O}(U)$，不计入。
见下面数学题单的「§1.5 因子」。
欢迎关注 B站@灵茶山艾府
我们可以遍历数组 nums 中的每个元素，依次判断这些元素是否恰好有四个因数。对于任一元素 x，我们可以用类似质数判定的方法得到它的因数个数，其本质为：如果整数 x 有因数 y，那么也必有因数 x/y，并且 y 和 x/y 中至少有一个不大于 sqrt(x)。这样我们只需要在 [1, sqrt(x)] 的区间内枚举可能为整数 x 的因数 y，并通过 x/y 得到整数 x 的其它因数，时间复杂度为 $O(\sqrt{x})$。
如果 x 恰好有四个因数，我们就将其因数之和累加到答案中。
###C++
class Solution {
public:
int sumFourDivisors(vector<int>& nums) {
int ans = 0;
for (int num: nums) {
// factor_cnt: 因数的个数
// factor_sum: 因数的和
int factor_cnt = 0, factor_sum = 0;
for (int i = 1; i * i <= num; ++i) {
if (num % i == 0) {
++factor_cnt;
factor_sum += i;
if (i * i != num) { // 判断 i 和 num/i 是否相等，若不相等才能将 num/i 看成新的因数
++factor_cnt;
factor_sum += num / i;
}
}
}
if (factor_cnt == 4) {
ans += factor_sum;
}
}
return ans;
}
};
###Java
class Solution {
public int sumFourDivisors(int[] nums) {
int ans = 0;
for (int num : nums) {
// factor_cnt: 因数的个数
// factor_sum: 因数的和
int factor_cnt = 0, factor_sum = 0;
for (int i = 1; i * i <= num; ++i) {
if (num % i == 0) {
++factor_cnt;
factor_sum += i;
if (i * i != num) { // 判断 i 和 num/i 是否相等，若不相等才能将 num/i 看成新的因数
++factor_cnt;
factor_sum += num / i;
}
}
}
if (factor_cnt == 4) {
ans += factor_sum;
}
}
return ans;
}
}
###Python
class Solution:
def sumFourDivisors(self, nums: List[int]) -> int:
ans = 0
for num in nums:
# factor_cnt: 因数的个数
# factor_sum: 因数的和
factor_cnt = factor_sum = 0
i = 1
while i * i <= num:
if num % i == 0:
factor_cnt += 1
factor_sum += i
if i * i != num: # 判断 i 和 num/i 是否相等，若不相等才能将 num/i 看成新的因数
factor_cnt += 1
factor_sum += num // i
i += 1
if factor_cnt == 4:
ans += factor_sum
return ans
###C#
public class Solution {
public int SumFourDivisors(int[] nums) {
int ans = 0;
foreach (int num in nums) {
// factor_cnt: 因数的个数
// factor_sum: 因数的和
int factor_cnt = 0, factor_sum = 0;
for (int i = 1; i * i <= num; ++i) {
if (num % i == 0) {
++factor_cnt;
factor_sum += i;
if (i * i != num) { // 判断 i 和 num/i 是否相等，若不相等才能将 num/i 看成新的因数
++factor_cnt;
factor_sum += num / i;
}
}
}
if (factor_cnt == 4) {
ans += factor_sum;
}
}
return ans;
}
}
###Go
func sumFourDivisors(nums []int) int {
ans := 0
for _, num := range nums {
// factor_cnt: 因数的个数
// factor_sum: 因数的和
factor_cnt, factor_sum := 0, 0
for i := 1; i*i <= num; i++ {
if num%i == 0 {
factor_cnt++
factor_sum += i
if i*i != num { // 判断 i 和 num/i 是否相等，若不相等才能将 num/i 看成新的因数
factor_cnt++
factor_sum += num / i
}
}
}
if factor_cnt == 4 {
ans += factor_sum
}
}
return ans
}
###C
int sumFourDivisors(int* nums, int numsSize) {
int ans = 0;
for (int idx = 0; idx < numsSize; idx++) {
int num = nums[idx];
// factor_cnt: 因数的个数
// factor_sum: 因数的和
int factor_cnt = 0, factor_sum = 0;
for (int i = 1; i * i <= num; ++i) {
if (num % i == 0) {
++factor_cnt;
factor_sum += i;
if (i * i != num) { // 判断 i 和 num/i 是否相等，若不相等才能将 num/i 看成新的因数
++factor_cnt;
factor_sum += num / i;
}
}
}
if (factor_cnt == 4) {
ans += factor_sum;
}
}
return ans;
}
###JavaScript
var sumFourDivisors = function(nums) {
let ans = 0;
for (const num of nums) {
// factor_cnt: 因数的个数
// factor_sum: 因数的和
let factor_cnt = 0, factor_sum = 0;
for (let i = 1; i * i <= num; ++i) {
if (num % i === 0) {
++factor_cnt;
factor_sum += i;
if (i * i !== num) { // 判断 i 和 num/i 是否相等，若不相等才能将 num/i 看成新的因数
++factor_cnt;
factor_sum += num / i;
}
}
}
if (factor_cnt === 4) {
ans += factor_sum;
}
}
return ans;
};
###TypeScript
function sumFourDivisors(nums: number[]): number {
let ans = 0;
for (const num of nums) {
// factor_cnt: 因数的个数
// factor_sum: 因数的和
let factor_cnt = 0, factor_sum = 0;
for (let i = 1; i * i <= num; ++i) {
if (num % i === 0) {
++factor_cnt;
factor_sum += i;
if (i * i !== num) { // 判断 i 和 num/i 是否相等，若不相等才能将 num/i 看成新的因数
++factor_cnt;
factor_sum += num / i;
}
}
}
if (factor_cnt === 4) {
ans += factor_sum;
}
}
return ans;
}
###Rust
impl Solution {
pub fn sum_four_divisors(nums: Vec<i32>) -> i32 {
let mut ans = 0;
for &num in &nums {
// factor_cnt: 因数的个数
// factor_sum: 因数的和
let mut factor_cnt = 0;
let mut factor_sum = 0;
let mut i = 1;
while i * i <= num {
if num % i == 0 {
factor_cnt += 1;
factor_sum += i;
if i * i != num { // 判断 i 和 num/i 是否相等，若不相等才能将 num/i 看成新的因数
factor_cnt += 1;
factor_sum += num / i;
}
}
i += 1;
}
if factor_cnt == 4 {
ans += factor_sum;
}
}
ans
}
}
复杂度分析
时间复杂度：$O(N\sqrt{C})$，其中 $N$ 是数组 nums 的长度，$C$ 是数组 nums 中元素值的范围，在本题中 $C$ 不超过 $10^5$。
空间复杂度：$O(1)$。
预备知识
分析与算法
直觉告诉我们，恰好有四个因数的整数不会有很多，我们是否可以预先找出它们呢？
根据「算数基本定理」（又叫「唯一分解定理」），如果整数 $x$ 可以分解为：
$$
x = p_1^{\alpha_1}p_2^{\alpha_2}\cdots p_k^{\alpha_k}
$$
其中 $p_i$ 为互不相同的质数（即 $x$ 的质因数）。那么 $x$ 的因数个数为：
$$
\textit{factor_count}(x) = \prod_{i=1}^k (\alpha_i + 1)
$$
如果 $\textit{factor_count}(x)$ 的值为 $4$，那么只有两种可能：
整数 $x$ 只有一个质因数，对应的指数为 $3$，此时 $\textit{factor_count}(x) = (3+1) = 4$；
整数 $x$ 有两个质因数，对应的指数均为 $1$，此时 $\textit{factor_count}(x) = (1+1)(1+1) = 4$。
对于第一种情况，我们需要找到所有不大于 $C^{1/3}$ 的质数；对于第二种情况，我们需要找到所有不大于 $C$ 的质数，再将它们两两相乘并筛去超过 $C$ 的那些结果。这里 $C$ 的定义与方法一中的复杂度分析部分一致。综上所述，我们需要找到所有不大于 $C$ 的质数。
我们如何找出所有不大于 $C$ 的质数呢？这时就需要「埃拉托斯特尼筛法」或「欧拉筛法」的帮助了。它们可以帮助我们快速找到这些质数。这两种筛法的算法细节不是这篇题解的重点，这里不再赘述。在找到了这些质数后，我们就可以构造出所有满足上述两种可能的 $x$ 了。我们将 $x$ 以及它的因数之和存入哈希映射（HashMap）中，这样就可以在 $O(1)$ 的时间判断数组 nums 中的每个元素是否满足要求，并统计满足要求的元素的因数之和了。
下面的代码给出了 Python 和 C++ 语言的「埃拉托斯特尼筛法」以及「欧拉筛法」的实现。
###C++
class Solution {
public:
int sumFourDivisors(vector<int>& nums) {
// C 是数组 nums 元素的上限，C3 是 C 的立方根
int C = 100000, C3 = 46;
vector<int> isprime(C + 1, 1);
vector<int> primes;
// 埃拉托斯特尼筛法
for (int i = 2; i <= C; ++i) {
if (isprime[i]) {
primes.push_back(i);
}
for (int j = i + i; j <= C; j += i) {
isprime[j] = 0;
}
}
// 欧拉筛法
/*
for (int i = 2; i <= C; ++i) {
if (isprime[i]) {
primes.push_back(i);
}
for (int prime: primes) {
if (i * prime > C) {
break;
}
isprime[i * prime] = 0;
if (i % prime == 0) {
break;
}
}
}
*/
// 通过质数表构造出所有的四因数
unordered_map<int, int> factor4;
for (int prime: primes) {
if (prime <= C3) {
factor4[prime * prime * prime] = 1 + prime + prime * prime + prime * prime * prime;
}
}
for (int i = 0; i < primes.size(); ++i) {
for (int j = i + 1; j < primes.size(); ++j) {
if (primes[i] <= C / primes[j]) {
factor4[primes[i] * primes[j]] = 1 + primes[i] + primes[j] + primes[i] * primes[j];
}
else {
break;
}
}
}
int ans = 0;
for (int num: nums) {
if (factor4.count(num)) {
ans += factor4[num];
}
}
return ans;
}
};
###Java
class Solution {
public int sumFourDivisors(int[] nums) {
// C 是数组 nums 元素的上限，C3 是 C 的立方根
int C = 100000, C3 = 46;
boolean[] isPrime = new boolean[C + 1];
Arrays.fill(isPrime, true);
List<Integer> primes = new ArrayList<Integer>();
// 埃拉托斯特尼筛法
for (int i = 2; i <= C; ++i) {
if (isPrime[i]) {
primes.add(i);
}
for (int j = i + i; j <= C; j += i) {
isPrime[j] = false;
}
}
// 欧拉筛法
/*
for (int i = 2; i <= C; ++i) {
if (isPrime[i]) {
primes.add(i);
}
for (int prime : primes) {
if (i * prime > C) {
break;
}
isPrime[i * prime] = false;
if (i % prime == 0) {
break;
}
}
}
*/
// 通过质数表构造出所有的四因数
Map<Integer, Integer> factor4 = new HashMap<Integer, Integer>();
for (int prime : primes) {
if (prime <= C3) {
factor4.put(prime * prime * prime, 1 + prime + prime * prime + prime * prime * prime);
}
}
for (int i = 0; i < primes.size(); ++i) {
for (int j = i + 1; j < primes.size(); ++j) {
if (primes.get(i) <= C / primes.get(j)) {
factor4.put(primes.get(i) * primes.get(j), 1 + primes.get(i) + primes.get(j) + primes.get(i) * primes.get(j));
} else {
break;
}
}
}
int ans = 0;
for (int num : nums) {
if (factor4.containsKey(num)) {
ans += factor4.get(num);
}
}
return ans;
}
}
###Python
class Solution:
def sumFourDivisors(self, nums: List[int]) -> int:
# C 是数组 nums 元素的上限，C3 是 C 的立方根
C, C3 = 100000, 46
isprime = [True] * (C + 1)
primes = list()
# 埃拉托斯特尼筛法
for i in range(2, C + 1):
if isprime[i]:
primes.append(i)
for j in range(i + i, C + 1, i):
isprime[j] = False
# 欧拉筛法
"""
for i in range(2, C + 1):
if isprime[i]:
primes.append(i)
for prime in primes:
if i * prime > C:
break
isprime[i * prime] = False
if i % prime == 0:
break
"""
# 通过质数表构造出所有的四因数
factor4 = dict()
for prime in primes:
if prime <= C3:
factor4[prime**3] = 1 + prime + prime**2 + prime**3
for i in range(len(primes)):
for j in range(i + 1, len(primes)):
if primes[i] * primes[j] <= C:
factor4[primes[i] * primes[j]] = 1 + primes[i] + primes[j] + primes[i] * primes[j]
else:
break
ans = 0
for num in nums:
if num in factor4:
ans += factor4[num]
return ans
###C#
public class Solution {
public int SumFourDivisors(int[] nums) {
// C 是数组 nums 元素的上限，C3 是 C 的立方根
int C = 100000, C3 = 46;
int[] isprime = new int[C + 1];
for (int i = 2; i <= C; i++) isprime[i] = 1;
List<int> primes = new List<int>();
// 埃拉托斯特尼筛法
for (int i = 2; i <= C; ++i) {
if (isprime[i] == 1) {
primes.Add(i);
}
for (int j = i + i; j <= C; j += i) {
isprime[j] = 0;
}
}
// 欧拉筛法
/*
for (int i = 2; i <= C; ++i) {
if (isprime[i] == 1) {
primes.Add(i);
}
foreach (int prime in primes) {
if (i * prime > C) {
break;
}
isprime[i * prime] = 0;
if (i % prime == 0) {
break;
}
}
}
*/
// 通过质数表构造出所有的四因数
Dictionary<int, int> factor4 = new Dictionary<int, int>();
foreach (int prime in primes) {
if (prime <= C3) {
factor4[prime * prime * prime] = 1 + prime + prime * prime + prime * prime * prime;
}
}
for (int i = 0; i < primes.Count; ++i) {
for (int j = i + 1; j < primes.Count; ++j) {
if (primes[i] <= C / primes[j]) {
factor4[primes[i] * primes[j]] = 1 + primes[i] + primes[j] + primes[i] * primes[j];
}
else {
break;
}
}
}
int ans = 0;
foreach (int num in nums) {
if (factor4.ContainsKey(num)) {
ans += factor4[num];
}
}
return ans;
}
}
###Go
func sumFourDivisors(nums []int) int {
// C 是数组 nums 元素的上限，C3 是 C 的立方根
C, C3 := 100000, 46
isprime := make([]int, C+1)
for i := 2; i <= C; i++ {
isprime[i] = 1
}
primes := []int{}
// 埃拉托斯特尼筛法
for i := 2; i <= C; i++ {
if isprime[i] == 1 {
primes = append(primes, i)
}
for j := i + i; j <= C; j += i {
isprime[j] = 0
}
}
// 欧拉筛法
/*
for i := 2; i <= C; i++ {
if isprime[i] == 1 {
primes = append(primes, i)
}
for _, prime := range primes {
if i * prime > C {
break
}
isprime[i * prime] = 0
if i % prime == 0 {
break
}
}
}
*/
// 通过质数表构造出所有的四因数
factor4 := make(map[int]int)
for _, prime := range primes {
if prime <= C3 {
factor4[prime * prime * prime] = 1 + prime + prime * prime + prime * prime * prime
}
}
for i := 0; i < len(primes); i++ {
for j := i + 1; j < len(primes); j++ {
if primes[i] <= C / primes[j] {
factor4[primes[i] * primes[j]] = 1 + primes[i] + primes[j] + primes[i] * primes[j]
} else {
break
}
}
}
ans := 0
for _, num := range nums {
if val, exists := factor4[num]; exists {
ans += val
}
}
return ans
}
###C
int sumFourDivisors(int* nums, int numsSize) {
// C 是数组 nums 元素的上限，C3 是 C 的立方根
const int C = 100000, C3 = 46;
int* isprime = (int*)malloc((C + 1) * sizeof(int));
memset(isprime, 0, (C + 1) * sizeof(int));
int* primes = (int*)malloc((C + 1) * sizeof(int));
int primeCount = 0;
// 埃拉托斯特尼筛法
for (int i = 2; i <= C; ++i) {
isprime[i] = 1;
}
for (int i = 2; i <= C; ++i) {
if (isprime[i]) {
primes[primeCount++] = i;
}
for (int j = i + i; j <= C; j += i) {
isprime[j] = 0;
}
}
// 欧拉筛法
/*
for (int i = 2; i <= C; ++i) {
if (isprime[i]) {
primes[primeCount++] = i;
}
for (int j = 0; j < primeCount; ++j) {
if (i * primes[j] > C) {
break;
}
isprime[i * primes[j]] = 0;
if (i % primes[j] == 0) {
break;
}
}
}
*/
// 通过质数表构造出所有的四因数
int* factor4_keys = (int*)malloc(primeCount * primeCount * sizeof(int));
int* factor4_values = (int*)malloc(primeCount * primeCount * sizeof(int));
int factor4_count = 0;
for (int i = 0; i < primeCount; ++i) {
int prime = primes[i];
if (prime <= C3) {
factor4_keys[factor4_count] = prime * prime * prime;
factor4_values[factor4_count] = 1 + prime + prime * prime + prime * prime * prime;
factor4_count++;
}
}
for (int i = 0; i < primeCount; ++i) {
for (int j = i + 1; j < primeCount; ++j) {
if (primes[i] <= C / primes[j]) {
factor4_keys[factor4_count] = primes[i] * primes[j];
factor4_values[factor4_count] = 1 + primes[i] + primes[j] + primes[i] * primes[j];
factor4_count++;
} else {
break;
}
}
}
int ans = 0;
for (int idx = 0; idx < numsSize; ++idx) {
int num = nums[idx];
for (int i = 0; i < factor4_count; ++i) {
if (factor4_keys[i] == num) {
ans += factor4_values[i];
break;
}
}
}
free(isprime);
free(primes);
free(factor4_keys);
free(factor4_values);
return ans;
}
###JavaScript
var sumFourDivisors = function(nums) {
// C 是数组 nums 元素的上限，C3 是 C 的立方根
const C = 100000, C3 = 46;
let isprime = new Array(C + 1).fill(0);
let primes = [];
// 埃拉托斯特尼筛法
for (let i = 2; i <= C; i++) {
isprime[i] = 1;
}
for (let i = 2; i <= C; i++) {
if (isprime[i]) {
primes.push(i);
}
for (let j = i + i; j <= C; j += i) {
isprime[j] = 0;
}
}
// 欧拉筛法
/*
for (let i = 2; i <= C; i++) {
if (isprime[i]) {
primes.push(i);
}
for (let prime of primes) {
if (i * prime > C) {
break;
}
isprime[i * prime] = 0;
if (i % prime === 0) {
break;
}
}
}
*/
// 通过质数表构造出所有的四因数
let factor4 = new Map();
for (let prime of primes) {
if (prime <= C3) {
factor4.set(prime * prime * prime, 1 + prime + prime * prime + prime * prime * prime);
}
}
for (let i = 0; i < primes.length; i++) {
for (let j = i + 1; j < primes.length; j++) {
if (primes[i] <= C / primes[j]) {
factor4.set(primes[i] * primes[j], 1 + primes[i] + primes[j] + primes[i] * primes[j]);
} else {
break;
}
}
}
let ans = 0;
for (let num of nums) {
if (factor4.has(num)) {
ans += factor4.get(num);
}
}
return ans;
};
###TypeScript
function sumFourDivisors(nums: number[]): number {
// C 是数组 nums 元素的上限，C3 是 C 的立方根
const C: number = 100000, C3: number = 46;
let isprime: number[] = new Array(C + 1).fill(0);
let primes: number[] = [];
// 埃拉托斯特尼筛法
for (let i = 2; i <= C; i++) {
isprime[i] = 1;
}
for (let i = 2; i <= C; i++) {
if (isprime[i]) {
primes.push(i);
}
for (let j = i + i; j <= C; j += i) {
isprime[j] = 0;
}
}
// 欧拉筛法
/*
for (let i = 2; i <= C; i++) {
if (isprime[i]) {
primes.push(i);
}
for (let prime of primes) {
if (i * prime > C) {
break;
}
isprime[i * prime] = 0;
if (i % prime === 0) {
break;
}
}
}
*/
// 通过质数表构造出所有的四因数
let factor4: Map<number, number> = new Map();
for (let prime of primes) {
if (prime <= C3) {
factor4.set(prime * prime * prime, 1 + prime + prime * prime + prime * prime * prime);
}
}
for (let i = 0; i < primes.length; i++) {
for (let j = i + 1; j < primes.length; j++) {
if (primes[i] <= C / primes[j]) {
factor4.set(primes[i] * primes[j], 1 + primes[i] + primes[j] + primes[i] * primes[j]);
} else {
break;
}
}
}
let ans: number = 0;
for (let num of nums) {
if (factor4.has(num)) {
ans += factor4.get(num)!;
}
}
return ans;
}
###Rust
use std::collections::HashMap;
impl Solution {
pub fn sum_four_divisors(nums: Vec<i32>) -> i32 {
// C 是数组 nums 元素的上限，C3 是 C 的立方根
const C: i32 = 100000;
const C3: i32 = 46;
let mut isprime = vec![0; (C + 1) as usize];
let mut primes = Vec::new();
// 埃拉托斯特尼筛法
for i in 2..=C {
isprime[i as usize] = 1;
}
for i in 2..=C {
if isprime[i as usize] == 1 {
primes.push(i);
}
let mut j = i + i;
while j <= C {
isprime[j as usize] = 0;
j += i;
}
}
// 欧拉筛法
/*
for i in 2..=C {
if isprime[i as usize] == 1 {
primes.push(i);
}
for &prime in &primes {
if i * prime > C {
break;
}
isprime[(i * prime) as usize] = 0;
if i % prime == 0 {
break;
}
}
}
*/
// 通过质数表构造出所有的四因数
let mut factor4 = HashMap::new();
for &prime in &primes {
if prime <= C3 {
let key = prime * prime * prime;
let value = 1 + prime + prime * prime + prime * prime * prime;
factor4.insert(key, value);
}
}
for i in 0..primes.len() {
for j in i + 1..primes.len() {
if primes[i] <= C / primes[j] {
let key = primes[i] * primes[j];
let value = 1 + primes[i] + primes[j] + primes[i] * primes[j];
factor4.insert(key, value);
} else {
break;
}
}
}
let mut ans = 0;
for num in nums {
if let Some(&value) = factor4.get(&num) {
ans += value;
}
}
ans
}
}
复杂度分析
时间复杂度：$O(\pi^2(C) + C\log\log C + N)$ 或 $O(\pi^2(C) + C + N)$，其中 $\pi(X)$ 为「质数计数函数」，表示不超过 $X$ 的质数个数。「埃拉托斯特尼筛法」的时间复杂度为 $O(C\log\log C)$，「欧拉筛法」的时间复杂度为 $O(C)$；通过质数表构造出所有四因数的时间复杂度为 $O(\pi(C^{1/3})) + O(\pi^2(C)) = O(\pi^2(C))$，遍历数组 nums 中的所有元素并检查是否为四因数的时间复杂度为 $O(N)$。
空间复杂度：$O(C + \pi(C))$，无论哪一种筛法，都需要长度为 $C$ 的数组记录每个数是否为质数，以及长度为 $\pi(C)$ 的数组存储所有的质数。
过不了的问题主要出在测试用的83521和2401这两个数上，分别是17和7的四次方。
写法没啥特别的，遍历到根号，计数，但是比赛全程都卡在最后这个巨长的用例上，始终找不到为啥。
这是比赛途中采用的代码：
#include <math.h>
class Solution {
public:
int sumFourDivisors(vector<int>& nums) {
int res = 0;
for(int i = 0; i < nums.size(); ++i){
int cnt = 0;
int sum = 0;
for(int j = 2; j < sqrt(nums[i]); ++j){
if(nums[i] % j == 0){
cnt++;
sum += j;
sum += nums[i] / j;
}
if(cnt > 1) break;
}
if(cnt == 1){
res = res + sum + nums[i] + 1;
}
}
return res;
}
};
这里我以为j < sqrt(nums[i])而不用等于就可以避免平方根的问题，但是显然忽略了更高阶的四次方情况。
排错方式比较智障，人工二分法测试，找到了83521和2401这两个数，发现其四次方特性。
参考评论区各位前辈代码，把平方判断单独拿出来写：
#include <math.h>
class Solution {
public:
int sumFourDivisors(vector<int>& nums) {
int res = 0;
for(int i = 0; i < nums.size(); ++i){
int cnt = 0;
int sum = 0;
int sq = sqrt(nums[i]);
for(int j = 2; j <= sq; ++j){
if(nums[i] % j == 0){
cnt++;
sum += j;
sum += nums[i] / j;
}
if(cnt > 1) break;
}
if(cnt == 1 && sq * sq != nums[i]){
res = res + sum + nums[i] + 1;
}
}
return res;
}
};
通过~
你有一个 n x 3 的网格图 grid ，你需要用 红，黄，绿 三种颜色之一给每一个格子上色，且确保相邻格子颜色不同（也就是有相同水平边或者垂直边的格子颜色不同）。
给你网格图的行数 n 。
请你返回给 grid 涂色的方案数。由于答案可能会非常大，请你返回答案对 10^9 + 7 取余的结果。
示例 1：
输入：n = 1 输出：12 解释：总共有 12 种可行的方法：
示例 2：
输入：n = 2 输出：54
示例 3：
输入：n = 3 输出：246
示例 4：
输入：n = 7 输出：106494
示例 5：
输入：n = 5000 输出：30228214
提示：
n == grid.lengthgrid[i].length == 31 <= n <= 5000考虑暴力搜索，枚举每个格子涂哪种颜色。
从下往上涂色（这样可以在多组测试数据间复用记忆化的结果），我们需要知道：
三种颜色分别用 $0,1,2$ 表示，存储一个格子的颜色信息需要 $2$ 个比特位，一行的颜色就需要 $2\cdot 3 = 6$ 个比特位。
注意取模。为什么可以在中途取模？原理见 模运算的世界：当加减乘除遇上取模。
###py
MOD = 1_000_000_007
# (i, j)：当前位置
# pre_row：上一行（i+1 行）的颜色
# cur_row：当前这一行已填入的颜色
@cache # 缓存装饰器，避免重复计算 dfs（一行代码实现记忆化）
def dfs(i: int, j: int, pre_row: int, cur_row: int) -> int:
if i == 0: # 所有格子都已涂色
return 1 # 找到一个合法方案
if j == 3: # i 行已涂色
# 开始对 i-1 行涂色，cur_row 变成 pre_row
return dfs(i - 1, 0, cur_row, 0)
res = 0
for color in range(3): # 枚举 (i, j) 的颜色 color
# 不能和下面相邻格子 (i+1, j) 颜色相同
# 不能和左侧相邻格子 (i, j-1) 颜色相同
if pre_row and color == pre_row >> (j * 2) & 3 or \
j and color == cur_row >> ((j - 1) * 2) & 3:
continue
res += dfs(i, j + 1, pre_row, cur_row | (color << (j * 2)))
return res % MOD
class Solution:
def numOfWays(self, n: int) -> int:
return dfs(n, 0, 0, 0) # 从最后一行开始涂色
###java
class Solution {
private static final int MOD = 1_000_000_007;
// 全局 memo，记忆化的内容可以在不同测试数据间共享
private static Map<Integer, Integer> memo = new HashMap<>();
public int numOfWays(int n) {
return dfs(n, 0, 0, 0);
}
// (i, j)：当前位置
// preRow：上一行（i+1 行）的颜色
// curRow：当前这一行已填入的颜色
private int dfs(int i, int j, int preRow, int curRow) {
if (i == 0) { // 所有格子都已涂色
return 1; // 找到一个合法方案
}
if (j == 3) { // i 行已涂色
// 开始对 i-1 行涂色，curRow 变成 preRow
return dfs(i - 1, 0, curRow, 0);
}
// 参数压缩到一个 int 中
int key = (i << 14) | (j << 12) | (preRow << 6) | curRow;
if (memo.containsKey(key)) { // 之前计算过
return memo.get(key);
}
int res = 0;
for (int color = 0; color < 3; color++) { // 枚举 (i, j) 的颜色 color
if (preRow > 0 && color == (preRow >> (j * 2) & 3) || // 不能和下面相邻格子 (i+1, j) 颜色相同
j > 0 && color == (curRow >> ((j - 1) * 2) & 3)) { // 不能和左侧相邻格子 (i, j-1) 颜色相同
continue;
}
res = (res + dfs(i, j + 1, preRow, curRow | (color << (j * 2)))) % MOD;
}
memo.put(key, res); // 记忆化
return res;
}
}
###cpp
constexpr int MOD = 1'000'000'007;
// 全局 memo，记忆化的内容可以在不同测试数据间共享
unordered_map<int, int> memo;
// (i, j)：当前位置
// pre_row：上一行（i+1 行）的颜色
// cur_row：当前这一行已填入的颜色
int dfs(int i, int j, int pre_row, int cur_row) {
if (i == 0) { // 所有格子都已涂色
return 1; // 找到一个合法方案
}
if (j == 3) { // i 行已涂色
// 开始对 i-1 行涂色，cur_row 变成 pre_row
return dfs(i - 1, 0, cur_row, 0);
}
// 参数压缩到一个 int 中
int key = (i << 14) | (j << 12) | (pre_row << 6) | cur_row;
if (memo.contains(key)) { // 之前计算过
return memo[key];
}
int res = 0;
for (int color = 0; color < 3; color++) { // 枚举 (i, j) 的颜色 color
if (pre_row > 0 && color == (pre_row >> (j * 2) & 3) || // 不能和下面相邻格子 (i+1, j) 颜色相同
j > 0 && color == (cur_row >> ((j - 1) * 2) & 3)) { // 不能和左侧相邻格子 (i, j-1) 颜色相同
continue;
}
res = (res + dfs(i, j + 1, pre_row, cur_row | (color << (j * 2)))) % MOD;
}
memo[key] = res; // 记忆化
return res;
}
class Solution {
public:
int numOfWays(int n) {
return dfs(n, 0, 0, 0);
}
};
###go
const mod = 1_000_000_007
type tuple struct{ i, j, preRow, curRow int }
// 全局 memo，记忆化的内容可以在不同测试数据间共享
var memo = map[tuple]int{}
// (i, j)：当前位置
// preRow：上一行（i+1 行）的颜色
// curRow：当前这一行已填入的颜色
func dfs(i, j, preRow, curRow int) int {
if i == 0 { // 所有格子都已涂色
return 1 // 找到一个合法方案
}
if j == 3 { // i 行已涂色
// 开始对 i-1 行涂色，curRow 变成 preRow
return dfs(i-1, 0, curRow, 0)
}
t := tuple{i, j, preRow, curRow}
if res, ok := memo[t]; ok { // 之前计算过
return res
}
res := 0
for color := range 3 { // 枚举 (i, j) 的颜色 color
if preRow > 0 && color == preRow>>(j*2)&3 || // 不能和下面相邻格子 (i+1, j) 颜色相同
j > 0 && color == curRow>>((j-1)*2)&3 { // 不能和左侧相邻格子 (i, j-1) 颜色相同
continue
}
res = (res + dfs(i, j+1, preRow, curRow|color<<(j*2))) % mod
}
memo[t] = res // 记忆化
return res
}
func numOfWays(n int) int {
return dfs(n, 0, 0, 0)
}
回顾方法一，总体上看，DP 过程是一个线性递推（从 $(n-1)\times 3$ 网格图最后一行的所有涂色方案线性转移到 $n\times 3$ 网格图的最后一行），所以必然有线性递推式。
定义 $f[n]$ 表示给 $n\times 3$ 网格图涂色的方案数。
根据 Berlekamp-Massey 算法，用方法一求出 $f$ 的连续若干项（具体多少在文章中有解释，本题只需 $4$ 项），就可以用 Berlekamp-Massey 算法直接得到线性递推式
$$
f[i] = 5\cdot f[i-1]-2\cdot f[i-2] \ \ (i \ge 3)
$$
初始值 $f[1] = 12,\ f[2] = 54$。
也可以倒推出 $f[0]=3$，把 $f[0]$ 和 $f[1]$ 作为初始值。
###py
class Solution:
def numOfWays(self, n: int) -> int:
MOD = 1_000_000_007
f = [0] * (n + 1)
f[0] = 3
f[1] = 12
for i in range(2, n + 1):
f[i] = (f[i - 1] * 5 - f[i - 2] * 2) % MOD
return f[n]
###java
class Solution {
public int numOfWays(int n) {
final int MOD = 1_000_000_007;
int[] f = new int[n + 1];
f[0] = 3;
f[1] = 12;
for (int i = 2; i <= n; i++) {
f[i] = (int) ((f[i - 1] * 5L - f[i - 2] * 2L) % MOD); // 注意这里有减法，结果可能是负数
}
return (f[n] + MOD) % MOD; // 保证结果非负
}
}
###cpp
class Solution {
public:
int numOfWays(int n) {
constexpr int MOD = 1'000'000'007;
vector<int> f(n + 1);
f[0] = 3;
f[1] = 12;
for (int i = 2; i <= n; i++) {
f[i] = (f[i - 1] * 5LL - f[i - 2] * 2LL) % MOD; // 注意这里有减法，结果可能是负数
}
return (f[n] + MOD) % MOD; // 保证结果非负
}
};
###go
func numOfWays(n int) int {
const mod = 1_000_000_007
f := make([]int, n+1)
f[0] = 3
f[1] = 12
for i := 2; i <= n; i++ {
f[i] = (f[i-1]*5 - f[i-2]*2) % mod // 注意这里有减法，结果可能是负数
}
return (f[n] + mod) % mod // 保证结果非负
}
把方法二的递推式用矩阵乘法表示，即
$$
\begin{bmatrix}
f[i] \
f[i-1] \
\end{bmatrix}
= \begin{bmatrix}
5 & -2 \
1 & 0 \
\end{bmatrix}
\begin{bmatrix}
f[i-1] \
f[i-2] \
\end{bmatrix}
$$
把上式中的三个矩阵分别记作 $F[i],M,F[i-1]$，即
$$
F[i] = M\times F[i-1]
$$
那么有
$$
\begin{aligned}
F[n] &= M\times F[n-1] \
&= M\times M\times F[n-2] \
&= M\times M\times M\times F[n-3] \
&\ \ \vdots \
&= M^{n-1}\times F[1] \
\end{aligned}
$$
其中 $M^{n-1}$ 可以用快速幂计算，原理请看【图解】一张图秒懂快速幂。
初始值
$$
F[1] = \begin{bmatrix}
f[1] \
f[0] \
\end{bmatrix}
= \begin{bmatrix}
12 \
3 \
\end{bmatrix}
$$
答案为 $f[n]$，即 $F[n]$ 的第一项。
###py
MOD = 1_000_000_007
# a @ b，其中 @ 是矩阵乘法
def mul(a: List[List[int]], b: List[List[int]]) -> List[List[int]]:
return [[sum(x * y for x, y in zip(row, col)) % MOD for col in zip(*b)]
for row in a]
# a^n @ f0
def pow_mul(a: List[List[int]], n: int, f0: List[List[int]]) -> List[List[int]]:
res = f0
while n:
if n & 1:
res = mul(a, res)
a = mul(a, a)
n >>= 1
return res
class Solution:
def numOfWays(self, n: int) -> int:
m = [[5, -2], [1, 0]]
f1 = [[12], [3]]
fn = pow_mul(m, n - 1, f1)
return fn[0][0]
###py
import numpy as np
class Solution:
def numOfWays(self, n: int) -> int:
MOD = 1_000_000_007
m = np.array([[5, -2], [1, 0]], dtype=object)
f1 = np.array([12, 3], dtype=object)
fn = np.linalg.matrix_power(m, n - 1) @ f1
return fn[0] % MOD
###java
class Solution {
private static final int MOD = 1_000_000_007;
public int numOfWays(int n) {
int[][] m = {
{5, -2},
{1, 0},
};
int[][] f1 = {{12}, {3}};
int[][] fn = powMul(m, n - 1, f1);
return (fn[0][0] + MOD) % MOD; // 保证结果非负
}
// a^n * f0
private int[][] powMul(int[][] a, int n, int[][] f0) {
int[][] res = f0;
while (n > 0) {
if ((n & 1) > 0) {
res = mul(a, res);
}
a = mul(a, a);
n >>= 1;
}
return res;
}
// 返回矩阵 a 和矩阵 b 相乘的结果
private int[][] mul(int[][] a, int[][] b) {
int[][] c = new int[a.length][b[0].length];
for (int i = 0; i < a.length; i++) {
for (int k = 0; k < a[i].length; k++) {
if (a[i][k] == 0) {
continue;
}
for (int j = 0; j < b[k].length; j++) {
c[i][j] = (int) ((c[i][j] + (long) a[i][k] * b[k][j]) % MOD);
}
}
}
return c;
}
}
###cpp
constexpr int MOD = 1'000'000'007;
using matrix = vector<vector<int>>;
// 返回矩阵 a 和矩阵 b 相乘的结果
matrix mul(matrix& a, matrix& b) {
int n = a.size(), m = b[0].size();
matrix c = matrix(n, vector<int>(m));
for (int i = 0; i < n; i++) {
for (int k = 0; k < a[i].size(); k++) {
if (a[i][k] == 0) {
continue;
}
for (int j = 0; j < m; j++) {
c[i][j] = (c[i][j] + 1LL * a[i][k] * b[k][j]) % MOD;
}
}
}
return c;
}
// a^n * f0
matrix pow_mul(matrix a, int n, matrix& f0) {
matrix res = f0;
while (n) {
if (n & 1) {
res = mul(a, res);
}
a = mul(a, a);
n >>= 1;
}
return res;
}
class Solution {
public:
int numOfWays(int n) {
matrix m = {
{5, -2},
{1, 0},
};
matrix f1 = {{12}, {3}};
matrix fn = pow_mul(m, n - 1, f1);
return (fn[0][0] + MOD) % MOD; // 保证结果非负
}
};
###go
const mod = 1_000_000_007
type matrix [][]int
func newMatrix(n, m int) matrix {
a := make(matrix, n)
for i := range a {
a[i] = make([]int, m)
}
return a
}
// 返回矩阵 a 和矩阵 b 相乘的结果
func (a matrix) mul(b matrix) matrix {
c := newMatrix(len(a), len(b[0]))
for i, row := range a {
for k, x := range row {
if x == 0 {
continue
}
for j, y := range b[k] {
c[i][j] = (c[i][j] + x*y) % mod
}
}
}
return c
}
// a^n * f0
func (a matrix) powMul(n int, f0 matrix) matrix {
res := f0
for ; n > 0; n /= 2 {
if n%2 > 0 {
res = a.mul(res)
}
a = a.mul(a)
}
return res
}
func numOfWays(n int) int {
m := matrix{
{5, -2},
{1, 0},
}
f1 := matrix{{12}, {3}}
fn := m.powMul(n-1, f1)
return (fn[0][0] + mod) % mod // 保证结果非负
}
注：也可以用 Kitamasa 算法 计算，在矩阵更大（递推式更长）时有优势。
见下面动态规划题单的「§9.5 轮廓线 DP」和「§11.6 矩阵快速幂优化 DP」。
欢迎关注 B站@灵茶山艾府
我们可以用 $f[i][\textit{type}]$ 表示当网格的大小为 $i \times 3$ 且最后一行的填色方法为 $\textit{type}$ 时的方案数。由于我们在填充第 $i$ 行时，会影响我们填充方案的只有它上面的那一行（即 $i - 1$ 行），因此用 $f[i][\textit{type}]$ 表示状态是合理的。
那么我们如何计算 $f[i][\textit{type}]$ 呢？可以发现：
首先，$\textit{type}$ 本身是要满足要求的。每一行有 $3$ 个网格，如果我们用 $0, 1, 2$ 分别代表红黄绿，那么 $\textit{type}$ 可以看成一个三进制数，例如 $\textit{type} = (102)_3$ 时，表示 $3$ 个网格从左到右的颜色分别为黄、红、绿；
其次，$f[i][\textit{type}]$ 应该等于所有 $f[i - 1][\textit{type}']$ 的和，其中 $\textit{type'}$ 和 $\textit{type}$ 可以作为相邻的行。也就是说，$\textit{type'}$ 和 $\textit{type}$ 的对应位置不能相同。
递推解法的本身不难想出，难度在于上述的预处理以及编码实现。下面给出包含详细注释的 C++、Java 和 Python 代码。
###C++
class Solution {
private:
static constexpr int mod = 1000000007;
public:
int numOfWays(int n) {
// 预处理出所有满足条件的 type
vector<int> types;
for (int i = 0; i < 3; ++i) {
for (int j = 0; j < 3; ++j) {
for (int k = 0; k < 3; ++k) {
if (i != j && j != k) {
// 只要相邻的颜色不相同就行
// 将其以十进制的形式存储
types.push_back(i * 9 + j * 3 + k);
}
}
}
}
int type_cnt = types.size();
// 预处理出所有可以作为相邻行的 type 对
vector<vector<int>> related(type_cnt, vector<int>(type_cnt));
for (int i = 0; i < type_cnt; ++i) {
// 得到 types[i] 三个位置的颜色
int x1 = types[i] / 9, x2 = types[i] / 3 % 3, x3 = types[i] % 3;
for (int j = 0; j < type_cnt; ++j) {
// 得到 types[j] 三个位置的颜色
int y1 = types[j] / 9, y2 = types[j] / 3 % 3, y3 = types[j] % 3;
// 对应位置不同色，才能作为相邻的行
if (x1 != y1 && x2 != y2 && x3 != y3) {
related[i][j] = 1;
}
}
}
// 递推数组
vector<vector<int>> f(n + 1, vector<int>(type_cnt));
// 边界情况，第一行可以使用任何 type
for (int i = 0; i < type_cnt; ++i) {
f[1][i] = 1;
}
for (int i = 2; i <= n; ++i) {
for (int j = 0; j < type_cnt; ++j) {
for (int k = 0; k < type_cnt; ++k) {
// f[i][j] 等于所有 f[i - 1][k] 的和
// 其中 k 和 j 可以作为相邻的行
if (related[k][j]) {
f[i][j] += f[i - 1][k];
f[i][j] %= mod;
}
}
}
}
// 最终所有的 f[n][...] 之和即为答案
int ans = 0;
for (int i = 0; i < type_cnt; ++i) {
ans += f[n][i];
ans %= mod;
}
return ans;
}
};
###Java
class Solution {
static final int MOD = 1000000007;
public int numOfWays(int n) {
// 预处理出所有满足条件的 type
List<Integer> types = new ArrayList<Integer>();
for (int i = 0; i < 3; ++i) {
for (int j = 0; j < 3; ++j) {
for (int k = 0; k < 3; ++k) {
if (i != j && j != k) {
// 只要相邻的颜色不相同就行
// 将其以十进制的形式存储
types.add(i * 9 + j * 3 + k);
}
}
}
}
int typeCnt = types.size();
// 预处理出所有可以作为相邻行的 type 对
int[][] related = new int[typeCnt][typeCnt];
for (int i = 0; i < typeCnt; ++i) {
// 得到 types[i] 三个位置的颜色
int x1 = types.get(i) / 9, x2 = types.get(i) / 3 % 3, x3 = types.get(i) % 3;
for (int j = 0; j < typeCnt; ++j) {
// 得到 types[j] 三个位置的颜色
int y1 = types.get(j) / 9, y2 = types.get(j) / 3 % 3, y3 = types.get(j) % 3;
// 对应位置不同色，才能作为相邻的行
if (x1 != y1 && x2 != y2 && x3 != y3) {
related[i][j] = 1;
}
}
}
// 递推数组
int[][] f = new int[n + 1][typeCnt];
// 边界情况，第一行可以使用任何 type
for (int i = 0; i < typeCnt; ++i) {
f[1][i] = 1;
}
for (int i = 2; i <= n; ++i) {
for (int j = 0; j < typeCnt; ++j) {
for (int k = 0; k < typeCnt; ++k) {
// f[i][j] 等于所有 f[i - 1][k] 的和
// 其中 k 和 j 可以作为相邻的行
if (related[k][j] != 0) {
f[i][j] += f[i - 1][k];
f[i][j] %= MOD;
}
}
}
}
// 最终所有的 f[n][...] 之和即为答案
int ans = 0;
for (int i = 0; i < typeCnt; ++i) {
ans += f[n][i];
ans %= MOD;
}
return ans;
}
}
###Python
class Solution:
def numOfWays(self, n: int) -> int:
mod = 10**9 + 7
# 预处理出所有满足条件的 type
types = list()
for i in range(3):
for j in range(3):
for k in range(3):
if i != j and j != k:
# 只要相邻的颜色不相同就行
# 将其以十进制的形式存储
types.append(i * 9 + j * 3 + k)
type_cnt = len(types)
# 预处理出所有可以作为相邻行的 type 对
related = [[0] * type_cnt for _ in range(type_cnt)]
for i, ti in enumerate(types):
# 得到 types[i] 三个位置的颜色
x1, x2, x3 = ti // 9, ti // 3 % 3, ti % 3
for j, tj in enumerate(types):
# 得到 types[j] 三个位置的颜色
y1, y2, y3 = tj // 9, tj // 3 % 3, tj % 3
# 对应位置不同色，才能作为相邻的行
if x1 != y1 and x2 != y2 and x3 != y3:
related[i][j] = 1
# 递推数组
f = [[0] * type_cnt for _ in range(n + 1)]
# 边界情况，第一行可以使用任何 type
f[1] = [1] * type_cnt
for i in range(2, n + 1):
for j in range(type_cnt):
for k in range(type_cnt):
# f[i][j] 等于所有 f[i - 1][k] 的和
# 其中 k 和 j 可以作为相邻的行
if related[k][j]:
f[i][j] += f[i - 1][k]
f[i][j] %= mod
# 最终所有的 f[n][...] 之和即为答案
ans = sum(f[n]) % mod
return ans
###C#
public class Solution {
private const int mod = 1000000007;
public int NumOfWays(int n) {
// 预处理出所有满足条件的 type
List<int> types = new List<int>();
for (int i = 0; i < 3; ++i) {
for (int j = 0; j < 3; ++j) {
for (int k = 0; k < 3; ++k) {
if (i != j && j != k) {
// 只要相邻的颜色不相同就行
// 将其以十进制的形式存储
types.Add(i * 9 + j * 3 + k);
}
}
}
}
int type_cnt = types.Count;
// 预处理出所有可以作为相邻行的 type 对
int[][] related = new int[type_cnt][];
for (int i = 0; i < type_cnt; ++i) {
related[i] = new int[type_cnt];
// 得到 types[i] 三个位置的颜色
int x1 = types[i] / 9, x2 = types[i] / 3 % 3, x3 = types[i] % 3;
for (int j = 0; j < type_cnt; ++j) {
// 得到 types[j] 三个位置的颜色
int y1 = types[j] / 9, y2 = types[j] / 3 % 3, y3 = types[j] % 3;
// 对应位置不同色，才能作为相邻的行
if (x1 != y1 && x2 != y2 && x3 != y3) {
related[i][j] = 1;
}
}
}
// 递推数组
int[][] f = new int[n + 1][];
for (int i = 0; i <= n; ++i) {
f[i] = new int[type_cnt];
}
// 边界情况，第一行可以使用任何 type
for (int i = 0; i < type_cnt; ++i) {
f[1][i] = 1;
}
for (int i = 2; i <= n; ++i) {
for (int j = 0; j < type_cnt; ++j) {
for (int k = 0; k < type_cnt; ++k) {
// f[i][j] 等于所有 f[i - 1][k] 的和
// 其中 k 和 j 可以作为相邻的行
if (related[k][j] == 1) {
f[i][j] = (f[i][j] + f[i - 1][k]) % mod;
}
}
}
}
// 最终所有的 f[n][...] 之和即为答案
int ans = 0;
for (int i = 0; i < type_cnt; ++i) {
ans = (ans + f[n][i]) % mod;
}
return ans;
}
}
###Go
func numOfWays(n int) int {
// 预处理出所有满足条件的 type
mod := 1000000007
types := []int{}
for i := 0; i < 3; i++ {
for j := 0; j < 3; j++ {
for k := 0; k < 3; k++ {
if i != j && j != k {
// 只要相邻的颜色不相同就行
// 将其以十进制的形式存储
types = append(types, i*9 + j*3 + k)
}
}
}
}
type_cnt := len(types)
// 预处理出所有可以作为相邻行的 type 对
related := make([][]int, type_cnt)
for i := range related {
related[i] = make([]int, type_cnt)
}
for i := 0; i < type_cnt; i++ {
// 得到 types[i] 三个位置的颜色
x1 := types[i] / 9
x2 := types[i] / 3 % 3
x3 := types[i] % 3
for j := 0; j < type_cnt; j++ {
// 得到 types[j] 三个位置的颜色
y1 := types[j] / 9
y2 := types[j] / 3 % 3
y3 := types[j] % 3
// 对应位置不同色，才能作为相邻的行
if x1 != y1 && x2 != y2 && x3 != y3 {
related[i][j] = 1
}
}
}
// 递推数组
f := make([][]int, n+1)
for i := range f {
f[i] = make([]int, type_cnt)
}
// 边界情况，第一行可以使用任何 type
for i := 0; i < type_cnt; i++ {
f[1][i] = 1
}
for i := 2; i <= n; i++ {
for j := 0; j < type_cnt; j++ {
for k := 0; k < type_cnt; k++ {
// f[i][j] 等于所有 f[i - 1][k] 的和
// 其中 k 和 j 可以作为相邻的行
if related[k][j] == 1 {
f[i][j] = (f[i][j] + f[i-1][k]) % mod
}
}
}
}
// 最终所有的 f[n][...] 之和即为答案
ans := 0
for i := 0; i < type_cnt; i++ {
ans = (ans + f[n][i]) % mod
}
return ans
}
###C
int numOfWays(int n) {
// 预处理出所有满足条件的 type
const int mod = 1000000007;
int types[12];
int type_cnt = 0;
for (int i = 0; i < 3; ++i) {
for (int j = 0; j < 3; ++j) {
for (int k = 0; k < 3; ++k) {
if (i != j && j != k) {
// 只要相邻的颜色不相同就行
// 将其以十进制的形式存储
types[type_cnt++] = i * 9 + j * 3 + k;
}
}
}
}
// 预处理出所有可以作为相邻行的 type 对
int related[12][12] = {0};
for (int i = 0; i < type_cnt; ++i) {
// 得到 types[i] 三个位置的颜色
int x1 = types[i] / 9, x2 = types[i] / 3 % 3, x3 = types[i] % 3;
for (int j = 0; j < type_cnt; ++j) {
// 得到 types[j] 三个位置的颜色
int y1 = types[j] / 9, y2 = types[j] / 3 % 3, y3 = types[j] % 3;
// 对应位置不同色，才能作为相邻的行
if (x1 != y1 && x2 != y2 && x3 != y3) {
related[i][j] = 1;
}
}
}
// 递推数组
int f[n + 1][type_cnt];
// 初始化
for (int i = 0; i <= n; ++i) {
for (int j = 0; j < type_cnt; ++j) {
f[i][j] = 0;
}
}
// 边界情况，第一行可以使用任何 type
for (int i = 0; i < type_cnt; ++i) {
f[1][i] = 1;
}
for (int i = 2; i <= n; ++i) {
for (int j = 0; j < type_cnt; ++j) {
for (int k = 0; k < type_cnt; ++k) {
// f[i][j] 等于所有 f[i - 1][k] 的和
// 其中 k 和 j 可以作为相邻的行
if (related[k][j]) {
f[i][j] = (f[i][j] + f[i - 1][k]) % mod;
}
}
}
}
// 最终所有的 f[n][...] 之和即为答案
int ans = 0;
for (int i = 0; i < type_cnt; ++i) {
ans = (ans + f[n][i]) % mod;
}
return ans;
}
###JavaScript
var numOfWays = function(n) {
// 预处理出所有满足条件的 type
const mod = 1000000007;
const types = [];
for (let i = 0; i < 3; ++i) {
for (let j = 0; j < 3; ++j) {
for (let k = 0; k < 3; ++k) {
if (i !== j && j !== k) {
// 只要相邻的颜色不相同就行
// 将其以十进制的形式存储
types.push(i * 9 + j * 3 + k);
}
}
}
}
const type_cnt = types.length;
// 预处理出所有可以作为相邻行的 type 对
const related = Array.from({length: type_cnt}, () => new Array(type_cnt).fill(0));
for (let i = 0; i < type_cnt; ++i) {
// 得到 types[i] 三个位置的颜色
const x1 = Math.floor(types[i] / 9);
const x2 = Math.floor(types[i] / 3) % 3;
const x3 = types[i] % 3;
for (let j = 0; j < type_cnt; ++j) {
// 得到 types[j] 三个位置的颜色
const y1 = Math.floor(types[j] / 9);
const y2 = Math.floor(types[j] / 3) % 3;
const y3 = types[j] % 3;
// 对应位置不同色，才能作为相邻的行
if (x1 !== y1 && x2 !== y2 && x3 !== y3) {
related[i][j] = 1;
}
}
}
// 递推数组
const f = Array.from({length: n + 1}, () => new Array(type_cnt).fill(0));
// 边界情况，第一行可以使用任何 type
for (let i = 0; i < type_cnt; ++i) {
f[1][i] = 1;
}
for (let i = 2; i <= n; ++i) {
for (let j = 0; j < type_cnt; ++j) {
for (let k = 0; k < type_cnt; ++k) {
// f[i][j] 等于所有 f[i - 1][k] 的和
// 其中 k 和 j 可以作为相邻的行
if (related[k][j]) {
f[i][j] = (f[i][j] + f[i - 1][k]) % mod;
}
}
}
}
// 最终所有的 f[n][...] 之和即为答案
let ans = 0;
for (let i = 0; i < type_cnt; ++i) {
ans = (ans + f[n][i]) % mod;
}
return ans;
};
###TypeScript
function numOfWays(n: number): number {
// 预处理出所有满足条件的 type
const mod: number = 1000000007;
const types: number[] = [];
for (let i = 0; i < 3; ++i) {
for (let j = 0; j < 3; ++j) {
for (let k = 0; k < 3; ++k) {
if (i !== j && j !== k) {
// 只要相邻的颜色不相同就行
// 将其以十进制的形式存储
types.push(i * 9 + j * 3 + k);
}
}
}
}
const type_cnt: number = types.length;
// 预处理出所有可以作为相邻行的 type 对
const related: number[][] = Array.from({length: type_cnt}, () => new Array(type_cnt).fill(0));
for (let i = 0; i < type_cnt; ++i) {
// 得到 types[i] 三个位置的颜色
const x1: number = Math.floor(types[i] / 9);
const x2: number = Math.floor(types[i] / 3) % 3;
const x3: number = types[i] % 3;
for (let j = 0; j < type_cnt; ++j) {
// 得到 types[j] 三个位置的颜色
const y1: number = Math.floor(types[j] / 9);
const y2: number = Math.floor(types[j] / 3) % 3;
const y3: number = types[j] % 3;
// 对应位置不同色，才能作为相邻的行
if (x1 !== y1 && x2 !== y2 && x3 !== y3) {
related[i][j] = 1;
}
}
}
// 递推数组
const f: number[][] = Array.from({length: n + 1}, () => new Array(type_cnt).fill(0));
// 边界情况，第一行可以使用任何 type
for (let i = 0; i < type_cnt; ++i) {
f[1][i] = 1;
}
for (let i = 2; i <= n; ++i) {
for (let j = 0; j < type_cnt; ++j) {
for (let k = 0; k < type_cnt; ++k) {
// f[i][j] 等于所有 f[i - 1][k] 的和
// 其中 k 和 j 可以作为相邻的行
if (related[k][j]) {
f[i][j] = (f[i][j] + f[i - 1][k]) % mod;
}
}
}
}
// 最终所有的 f[n][...] 之和即为答案
let ans: number = 0;
for (let i = 0; i < type_cnt; ++i) {
ans = (ans + f[n][i]) % mod;
}
return ans;
}
###Rust
impl Solution {
pub fn num_of_ways(n: i32) -> i32 {
// 预处理出所有满足条件的 type
let mod_val = 1000000007;
let n = n as usize;
let mut types = Vec::new();
for i in 0..3 {
for j in 0..3 {
for k in 0..3 {
if i != j && j != k {
// 只要相邻的颜色不相同就行
// 将其以十进制的形式存储
types.push(i * 9 + j * 3 + k);
}
}
}
}
let type_cnt = types.len();
// 预处理出所有可以作为相邻行的 type 对
let mut related = vec![vec![0; type_cnt]; type_cnt];
for i in 0..type_cnt {
// 得到 types[i] 三个位置的颜色
let x1 = types[i] / 9;
let x2 = types[i] / 3 % 3;
let x3 = types[i] % 3;
for j in 0..type_cnt {
// 得到 types[j] 三个位置的颜色
let y1 = types[j] / 9;
let y2 = types[j] / 3 % 3;
let y3 = types[j] % 3;
// 对应位置不同色，才能作为相邻的行
if x1 != y1 && x2 != y2 && x3 != y3 {
related[i][j] = 1;
}
}
}
// 递推数组
let mut f = vec![vec![0; type_cnt]; n + 1];
// 边界情况，第一行可以使用任何 type
for i in 0..type_cnt {
f[1][i] = 1;
}
for i in 2..=n {
for j in 0..type_cnt {
for k in 0..type_cnt {
// f[i][j] 等于所有 f[i - 1][k] 的和
// 其中 k 和 j 可以作为相邻的行
if related[k][j] == 1 {
f[i][j] = (f[i][j] + f[i - 1][k]) % mod_val;
}
}
}
}
// 最终所有的 f[n][...] 之和即为答案
let mut ans = 0;
for i in 0..type_cnt {
ans = (ans + f[n][i]) % mod_val;
}
ans
}
}
复杂度分析
时间复杂度：$O(T^2N)$，其中 $T$ 是满足要求的 $\textit{type}$ 的数量，在示例一中已经给出了 $T = 12$。在递推的过程中，我们需要计算所有的 $f[i][\textit{type}]$，并且需要枚举上一行的 $\textit{type}'$。
空间复杂度：$O(T^2 + TN)$。我们需要 $T * T$ 的二维数组存储 $\textit{type}$ 之间的关系，$T * N$ 的数组存储递推的结果。注意到由于 $f[i][\textit{type}]$ 只和上一行的状态有关，我们可以使用两个一维数组存储当前行和上一行的 $f$ 值，空间复杂度降低至 $O(T^2 + 2T) = O(T^2)$。
如果读者有一些高中数学竞赛基础，就可以发现上面的这个递推式是线性的，也就是说：
我们可以进行一些化简；
它存在通项公式。
直观上，我们怎么化简方法一中的递推呢？
我们把满足要求的 $\textit{type}$ 都写出来，一共有 $12$ 种：
010, 012, 020, 021, 101, 102, 120, 121, 201, 202, 210, 212
我们可以把它们分成两类：
ABC 类：三个颜色互不相同，一共有 $6$ 种：012, 021, 102, 120, 201, 210；
ABA 类：左右两侧的颜色相同，也有 $6$ 种：010, 020, 101, 121, 202, 212。
这样我们就可以把 $12$ 种 $\textit{type}$ 浓缩成了 $2$ 种，尝试写出这两类之间的递推式。我们用 $f[i][0]$ 表示 ABC 类，$f[i][1]$ 表示 ABA 类。在计算时，我们可以将任意一种满足要求的涂色方法带入第 i - 1 行，并检查第 i 行的方案数，这是因为同一类的涂色方法都是等价的：
第 i - 1 行是 ABC 类，第 i 行是 ABC 类：以 012 为例，那么第 i 行只能是120 或 201，方案数为 $2$；
第 i - 1 行是 ABC 类，第 i 行是 ABA 类：以 012 为例，那么第 i 行只能是 101 或 121，方案数为 $2$；
第 i - 1 行是 ABA 类，第 i 行是 ABC 类：以 010 为例，那么第 i 行只能是 102 或 201，方案数为 2；
第 i - 1 行是 ABA 类，第 i 行是 ABA 类：以 010 为例，那么第 i 行只能是 101，121 或 202，方案数为 3。
因此我们就可以写出递推式：
$$
\begin{cases}
f[i][0] = 2 * f[i - 1][0] + 2 * f[i - 1][1] \
f[i][1] = 2 * f[i - 1][0] + 3 * f[i - 1][1]
\end{cases}
$$
###C++
class Solution {
private:
static constexpr int mod = 1000000007;
public:
int numOfWays(int n) {
int fi0 = 6, fi1 = 6;
for (int i = 2; i <= n; ++i) {
int new_fi0 = (2LL * fi0 + 2LL * fi1) % mod;
int new_fi1 = (2LL * fi0 + 3LL * fi1) % mod;
fi0 = new_fi0;
fi1 = new_fi1;
}
return (fi0 + fi1) % mod;
}
};
###Java
class Solution {
static final int MOD = 1000000007;
public int numOfWays(int n) {
long fi0 = 6, fi1 = 6;
for (int i = 2; i <= n; ++i) {
long newFi0 = (2 * fi0 + 2 * fi1) % MOD;
long newFi1 = (2 * fi0 + 3 * fi1) % MOD;
fi0 = newFi0;
fi1 = newFi1;
}
return (int) ((fi0 + fi1) % MOD);
}
}
###Python
class Solution:
def numOfWays(self, n: int) -> int:
mod = 10**9 + 7
fi0, fi1 = 6, 6
for i in range(2, n + 1):
fi0, fi1 = (2 * fi0 + 2 * fi1) % mod, (2 * fi0 + 3 * fi1) % mod
return (fi0 + fi1) % mod
###C#
public class Solution {
private const int mod = 1000000007;
public int NumOfWays(int n) {
long fi0 = 6, fi1 = 6;
for (int i = 2; i <= n; ++i) {
long new_fi0 = (2 * fi0 + 2 * fi1) % mod;
long new_fi1 = (2 * fi0 + 3 * fi1) % mod;
fi0 = new_fi0;
fi1 = new_fi1;
}
return (int)((fi0 + fi1) % mod);
}
}
###Go
func numOfWays(n int) int {
mod := 1000000007
fi0, fi1 := 6, 6
for i := 2; i <= n; i++ {
new_fi0 := (2*fi0 + 2*fi1) % mod
new_fi1 := (2*fi0 + 3*fi1) % mod
fi0, fi1 = new_fi0, new_fi1
}
return (fi0 + fi1) % mod
}
###C
int numOfWays(int n) {
const int mod = 1000000007;
long long fi0 = 6, fi1 = 6;
for (int i = 2; i <= n; ++i) {
long long new_fi0 = (2 * fi0 + 2 * fi1) % mod;
long long new_fi1 = (2 * fi0 + 3 * fi1) % mod;
fi0 = new_fi0;
fi1 = new_fi1;
}
return (fi0 + fi1) % mod;
}
###JavaScript
var numOfWays = function(n) {
const mod = 1000000007;
let fi0 = 6, fi1 = 6;
for (let i = 2; i <= n; i++) {
const new_fi0 = (2 * fi0 + 2 * fi1) % mod;
const new_fi1 = (2 * fi0 + 3 * fi1) % mod;
fi0 = new_fi0;
fi1 = new_fi1;
}
return (fi0 + fi1) % mod;
};
###TypeScript
function numOfWays(n: number): number {
const mod: number = 1000000007;
let fi0: number = 6, fi1: number = 6;
for (let i = 2; i <= n; i++) {
const new_fi0: number = (2 * fi0 + 2 * fi1) % mod;
const new_fi1: number = (2 * fi0 + 3 * fi1) % mod;
fi0 = new_fi0;
fi1 = new_fi1;
}
return (fi0 + fi1) % mod;
}
###Rust
impl Solution {
pub fn num_of_ways(n: i32) -> i32 {
let mod_val: i64 = 1000000007;
let mut fi0: i64 = 6;
let mut fi1: i64 = 6;
for _ in 2..= n {
let new_fi0 = (2 * fi0 + 2 * fi1) % mod_val;
let new_fi1 = (2 * fi0 + 3 * fi1) % mod_val;
fi0 = new_fi0;
fi1 = new_fi1;
}
((fi0 + fi1) % mod_val) as i32
}
}
复杂度分析
时间复杂度：$O(N)$。
空间复杂度：$O(1)$。
