func getMaxGridHappiness(m int, n int, introvertsCount int, extrovertsCount int) int {
p := int(math.Pow(3, float64(n-1)))
h := [3][3]int{{0, 0, 0}, {0, -60, -10}, {0, -10, 40}}
memo := make([][][][]int, m*n)
for i := range memo {
memo[i] = make([][][]int, p*3)
for j := range memo[i] {
memo[i][j] = make([][]int, introvertsCount+1)
for k := range memo[i][j] {
memo[i][j][k] = make([]int, extrovertsCount+1)
for l := range memo[i][j][k] {
memo[i][j][k][l] = -1
}
}
}
}
var dfs func(int, int, int, int) int
dfs = func(pos, pre, ic, ec int) int {
if pos == m*n || (ic == 0 && ec == 0) {
return 0
}
if memo[pos][pre][ic][ec] != -1 {
return memo[pos][pre][ic][ec]
}
ans := 0
up := pre / p
left := pre % 3
if pos%n == 0 {
left = 0
}
for i := 0; i < 3; i++ {
if (i == 1 && ic == 0) || (i == 2 && ec == 0) {
continue
}
cur := pre%p*3 + i
nic, nec := ic, ec
c := 0
if i == 1 {
nic--
c = 120
} else if i == 2 {
nec--
c = 40
}
a := h[up][i] + h[left][i]
b := dfs(pos+1, cur, nic, nec)
ans = max(ans, a+b+c)
}
memo[pos][pre][ic][ec] = ans
return ans
}
return dfs(0, 0, introvertsCount, extrovertsCount)
}
func max(a, b int) int {
if a > b {
return a
}
return b
}
class Solution {
private int m;
private int n;
private int p;
private final int[][] h = {{0, 0, 0}, {0, -60, -10}, {0, -10, 40}};
private Integer[][][][] memo;
public int getMaxGridHappiness(int m, int n, int introvertsCount, int extrovertsCount) {
this.m = m;
this.n = n;
p = (int) Math.pow(3, n - 1);
memo = new Integer[m * n][p * 3][introvertsCount + 1][extrovertsCount + 1];
return dfs(0, 0, introvertsCount, extrovertsCount);
}
private int dfs(int pos, int pre, int ic, int ec) {
if (pos == m * n || (ic == 0 && ec == 0)) {
return 0;
}
if (memo[pos][pre][ic][ec] != null) {
return memo[pos][pre][ic][ec];
}
int ans = 0;
int up = pre / p;
int left = pos % n == 0 ? 0 : pre % 3;
for (int i = 0; i < 3; ++i) {
if (i == 1 && (ic == 0) || (i == 2 && ec == 0)) {
continue;
}
int cur = pre % p * 3 + i;
int a = h[up][i] + h[left][i];
int b = dfs(pos + 1, cur, ic - (i == 1 ? 1 : 0), ec - (i == 2 ? 1 : 0));
int c = i == 1 ? 120 : (i == 2 ? 40 : 0);
ans = Math.max(ans, a + b + c);
}
return memo[pos][pre][ic][ec] = ans;
}
}
class Solution:
def getMaxGridHappiness(
self, m: int, n: int, introvertsCount: int, extrovertsCount: int
) -> int:
@cache
def dfs(pos: int, pre: int, ic: int, ec: int) -> int:
if pos == m * n or (ic == 0 and ec == 0):
return 0
ans = 0
up = pre // p
left = 0 if pos % n == 0 else pre % 3
for i in range(3):
if (i == 1 and ic == 0) or (i == 2 and ec == 0):
continue
cur = pre % p * 3 + i
a = h[up][i] + h[left][i]
b = dfs(pos + 1, cur, ic - (i == 1), ec - (i == 2))
c = 0
if i == 1:
c = 120
elif i == 2:
c = 40
ans = max(ans, a + b + c)
return ans
p = pow(3, n - 1)
h = [[0, 0, 0], [0, -60, -10], [0, -10, 40]]
return dfs(0, 0, introvertsCount, extrovertsCount)
func getMaxGridHappiness(m int, n int, introvertsCount int, extrovertsCount int) int {
mx := int(math.Pow(3, float64(n)))
f := make([]int, mx)
g := make([][]int, mx)
h := [3][3]int{{0, 0, 0}, {0, -60, -10}, {0, -10, 40}}
bits := make([][]int, mx)
ix := make([]int, mx)
ex := make([]int, mx)
memo := make([][][][]int, m)
for i := range g {
g[i] = make([]int, mx)
bits[i] = make([]int, n)
}
for i := range memo {
memo[i] = make([][][]int, mx)
for j := range memo[i] {
memo[i][j] = make([][]int, introvertsCount+1)
for k := range memo[i][j] {
memo[i][j][k] = make([]int, extrovertsCount+1)
for l := range memo[i][j][k] {
memo[i][j][k][l] = -1
}
}
}
}
for i := 0; i < mx; i++ {
mask := i
for j := 0; j < n; j++ {
x := mask % 3
mask /= 3
bits[i][j] = x
if x == 1 {
ix[i]++
f[i] += 120
} else if x == 2 {
ex[i]++
f[i] += 40
}
if j > 0 {
f[i] += h[x][bits[i][j-1]]
}
}
}
for i := 0; i < mx; i++ {
for j := 0; j < mx; j++ {
for k := 0; k < n; k++ {
g[i][j] += h[bits[i][k]][bits[j][k]]
}
}
}
var dfs func(int, int, int, int) int
dfs = func(i, pre, ic, ec int) int {
if i == m || (ic == 0 && ec == 0) {
return 0
}
if memo[i][pre][ic][ec] != -1 {
return memo[i][pre][ic][ec]
}
ans := 0
for cur := 0; cur < mx; cur++ {
if ix[cur] <= ic && ex[cur] <= ec {
ans = max(ans, f[cur]+g[pre][cur]+dfs(i+1, cur, ic-ix[cur], ec-ex[cur]))
}
}
memo[i][pre][ic][ec] = ans
return ans
}
return dfs(0, 0, introvertsCount, extrovertsCount)
}
func max(a, b int) int {
if a > b {
return a
}
return b
}
'''
状态压缩dp
每个网络三种状态：空置、内向、外向，记为0,1,2.
对每一行的状态使用长度为n的三进制数进行编码。
d(row, premask, iv, ev) 表示第row-1行的状态为premask,
并且row ~ m-1行最多放置iv个内向和ev个外向的人时的分数。
'''
class Solution:
def getMaxGridHappiness(self, m: int, n: int, introvertsCount: int, extrovertsCount: int) -> int:
T, N, M = 243, 5, 6
# 邻居间的分数
score = [
[0, 0, 0],
[0, -60, -10],
[0, -10, 40],
]
tot = 3**n
mask_bits = [[0] * N for _ in range(T)]
iv_count, ev_count = [0] * T, [0] * T
inner_score = [0] * T
inter_score = [[0] * T for _ in range(T)]
def init_data() -> None:
# 计算行内分数
for mask in range(tot):
mask_tmp = mask
for i in range(n):
x = mask_tmp % 3
mask_bits[mask][i] = x
mask_tmp //= 3
if x == 1:
iv_count[mask] += 1
inner_score[mask] += 120
elif x == 2:
ev_count[mask] += 1
inner_score[mask] += 40
if i > 0:
inner_score[mask] += score[x][mask_bits[mask][i - 1]]
# 计算行间分数
for i in range(tot):
for j in range(tot):
for k in range(n):
inter_score[i][j] += score[mask_bits[i][k]][mask_bits[j][k]]
@cache
def dfs(row: int, premask: int, iv: int, ev: int) -> int:
if row == m or (iv == 0 and ev == 0):
return 0
res = 0
for mask in range(tot):
# mask 包含的内向人数不能超过 iv ，外向人数不能超过 ev
if iv_count[mask] > iv or ev_count[mask] > ev:
continue
res = max(res, dfs(row + 1, mask, iv - iv_count[mask], ev - ev_count[mask]) + inner_score[mask] + inter_score[premask][mask])
return res
init_data()
return dfs(0, 0, introvertsCount, extrovertsCount)
