1659. Maximize Grid Happiness

Problem

You are given four integers, m, n, introvertsCount, and extrovertsCount. You have an m x n grid, and there are two types of people: introverts and extroverts. There are introvertsCount introverts and extrovertsCount extroverts.

You should decide how many people you want to live in the grid and assign each of them one grid cell. Note that you do not have to have all the people living in the grid.

The happiness of each person is calculated as follows:

Introverts start with 120 happiness and lose 30 happiness for each neighbor (introvert or extrovert).
Extroverts start with 40 happiness and gain 20 happiness for each neighbor (introvert or extrovert).

Neighbors live in the directly adjacent cells north, east, south, and west of a person's cell.

The grid happiness is the sum of each person's happiness. Return the maximum possible grid happiness.

Example 1:

Input: m = 2, n = 3, introvertsCount = 1, extrovertsCount = 2
Output: 240
Explanation: Assume the grid is 1-indexed with coordinates (row, column).
We can put the introvert in cell (1,1) and put the extroverts in cells (1,3) and (2,3).
- Introvert at (1,1) happiness: 120 (starting happiness) - (0 * 30) (0 neighbors) = 120
- Extrovert at (1,3) happiness: 40 (starting happiness) + (1 * 20) (1 neighbor) = 60
- Extrovert at (2,3) happiness: 40 (starting happiness) + (1 * 20) (1 neighbor) = 60
The grid happiness is 120 + 60 + 60 = 240.
The above figure shows the grid in this example with each person's happiness. The introvert stays in the light green cell while the extroverts live on the light purple cells.

Example 2:

Input: m = 3, n = 1, introvertsCount = 2, extrovertsCount = 1
Output: 260
Explanation: Place the two introverts in (1,1) and (3,1) and the extrovert at (2,1).
- Introvert at (1,1) happiness: 120 (starting happiness) - (1 * 30) (1 neighbor) = 90
- Extrovert at (2,1) happiness: 40 (starting happiness) + (2 * 20) (2 neighbors) = 80
- Introvert at (3,1) happiness: 120 (starting happiness) - (1 * 30) (1 neighbor) = 90
The grid happiness is 90 + 80 + 90 = 260.

Example 3:

Input: m = 2, n = 2, introvertsCount = 4, extrovertsCount = 0
Output: 240

Constraints:

1 <= m, n <= 5
0 <= introvertsCount, extrovertsCount <= min(m * n, 6)

Problem Summary

Given four integers m, n, introvertsCount, and extrovertsCount. There is an m x n grid and two types of people: introverts and extroverts. There are introvertsCount introverts and extrovertsCount extroverts in total. You need to decide how many people should live in the grid and assign each person to a grid cell. Note that you do not have to make all people live in the grid. Each person's happiness is calculated as follows:

An introvert starts with 120 happiness, but for every neighbor (introvert or extrovert), they lose 30 happiness.
An extrovert starts with 40 happiness, and for every neighbor (introvert or extrovert), they gain 20 happiness.

Neighbors are other people living in the four directly adjacent cells above, below, left, and right of a person's cell. Grid happiness is the sum of each person's happiness. Return the maximum possible grid happiness.

Solution Approach

Given an m x n grid and two types of people, determine how to arrange these two types of people to maximize the grid score. The two types of people have their own initial scores, and adjacency may either add or subtract points.
This problem has many states. First, each cell has 3 states, so each row has 3^6 = 729 different states. The score change within a row may be -60 (two introverts), +40 (two extroverts), or -10 (one introvert and one extrovert). The score change between two rows may be -60 (two introverts), +40 (two extroverts), or -10 (one introvert and one extrovert). Therefore, we can compress each row's state into a ternary number, so the grid becomes one-dimensional. The relationship between every two ternary numbers is the inter-row relationship, and within each ternary number we still need to determine the final intra-row score based on the number of introverts and extroverts. Define dp[lineStatusLast][row][introvertsCount][extrovertsCount] as the maximum score obtainable when the previous row row - 1 has state lineStatusLast, the current enumeration has reached row row, and there are still introvertsCount introverts and extrovertsCount extroverts left. The state transition equation is dp[lineStatusLast(row-1)][row][introvertsCount][extrovertsCount] = max{dp[lineStatusLast(row)][row+1][introvertsCount - countIC(lineStatusLast(row)) ][extrovertsCount - countEC(lineStatusLast(row)) ] + scoreInner(lineStatusLast(row)) + scoreOuter(lineStatusLast(row-1),lineStatusLast(row))}. There are 2 counting functions here: countIC counts how many introverts are in the ternary representation of the current row state. countEC counts how many extroverts are in the ternary representation of the current row state. scoreInner calculates the intra-row score of the ternary representation of the current row state. scoreOuter calculates the inter-row score between row row -1 and row row.
Since the computation required by this state transition equation is huge, some calculation results need to be precomputed. For example, compute the intra-row and inter-row scores corresponding to the 729 row states in advance. During dynamic programming state transitions, the scores can then be obtained directly from a lookup table. In this way, during DFS, using memoization for dp can greatly reduce the time complexity.
The problem also mentions that not all people need to be used. If introvertsCount = 0, extrovertsCount = 0, meaning all people have been used, then dp = 0. If row = m, meaning all rows have already been enumerated, then no matter how many people remain, the dp for this row is 0.
During initialization, note that the case of 0 should be handled specially, and row 0 and column 0 are both initialized to -1.

Code

package leetcode

import (
	"math"
)

func getMaxGridHappiness(m int, n int, introvertsCount int, extrovertsCount int) int {
	// lineStatus 					Encode the 3 states in each row: empty - 0, introvert - 1, extrovert - 2; each row state is represented in ternary
	// lineStatusList[729][6] 		The ternary representation of each row
	// introvertsCountInner[729] 	The number of introverts contained in each lineStatus
	// extrovertsCountInner[729] 	The number of extroverts contained in each lineStatus
	// scoreInner[729] 			 	The intra-row score contained in each lineStatus (only counts the score of lineStatus itself, not including its score with the previous row)
	// scoreOuter[729][729] 	 	The inter-row score contained in each lineStatus
	// dp[previous row's lineStatus][current row being processed][remaining introverts][remaining extroverts]
	n3, lineStatus, introvertsCountInner, extrovertsCountInner, scoreInner, scoreOuter, lineStatusList, dp := math.Pow(3.0, float64(n)), 0, [729]int{}, [729]int{}, [729]int{}, [729][729]int{}, [729][6]int{}, [729][6][7][7]int{}
	for i := 0; i < 729; i++ {
		lineStatusList[i] = [6]int{}
	}
	for i := 0; i < 729; i++ {
		dp[i] = [6][7][7]int{}
		for j := 0; j < 6; j++ {
			dp[i][j] = [7][7]int{}
			for k := 0; k < 7; k++ {
				dp[i][j][k] = [7]int{-1, -1, -1, -1, -1, -1, -1}
			}
		}
	}
	// Preprocessing
	for lineStatus = 0; lineStatus < int(n3); lineStatus++ {
		tmp := lineStatus
		for i := 0; i < n; i++ {
			lineStatusList[lineStatus][i] = tmp % 3
			tmp /= 3
		}
		introvertsCountInner[lineStatus], extrovertsCountInner[lineStatus], scoreInner[lineStatus] = 0, 0, 0
		for i := 0; i < n; i++ {
			if lineStatusList[lineStatus][i] != 0 {
				// Individual score
				if lineStatusList[lineStatus][i] == 1 {
					introvertsCountInner[lineStatus]++
					scoreInner[lineStatus] += 120
				} else if lineStatusList[lineStatus][i] == 2 {
					extrovertsCountInner[lineStatus]++
					scoreInner[lineStatus] += 40
				}
				// Intra-row score
				if i-1 >= 0 {
					scoreInner[lineStatus] += closeScore(lineStatusList[lineStatus][i], lineStatusList[lineStatus][i-1])
				}
			}
		}
	}
	// Inter-row score
	for lineStatus0 := 0; lineStatus0 < int(n3); lineStatus0++ {
		for lineStatus1 := 0; lineStatus1 < int(n3); lineStatus1++ {
			scoreOuter[lineStatus0][lineStatus1] = 0
			for i := 0; i < n; i++ {
				scoreOuter[lineStatus0][lineStatus1] += closeScore(lineStatusList[lineStatus0][i], lineStatusList[lineStatus1][i])
			}
		}
	}
	return dfs(0, 0, introvertsCount, extrovertsCount, m, int(n3), &dp, &introvertsCountInner, &extrovertsCountInner, &scoreInner, &scoreOuter)
}

// Score to add if x and y are adjacent
func closeScore(x, y int) int {
	if x == 0 || y == 0 {
		return 0
	}
	// Two introverts: each loses -30, total -60
	if x == 1 && y == 1 {
		return -60
	}
	if x == 2 && y == 2 {
		return 40
	}
	return -10
}

// dfs(previous row's lineStatus, current row being processed, remaining introverts, remaining extroverts)
func dfs(lineStatusLast, row, introvertsCount, extrovertsCount, m, n3 int, dp *[729][6][7][7]int, introvertsCountInner, extrovertsCountInner, scoreInner *[729]int, scoreOuter *[729][729]int) int {
	// Boundary condition: if processing is complete, or there are no people left
	if row == m || introvertsCount+extrovertsCount == 0 {
		return 0
	}
	// Memoization
	if dp[lineStatusLast][row][introvertsCount][extrovertsCount] != -1 {
		return dp[lineStatusLast][row][introvertsCount][extrovertsCount]
	}
	best := 0
	for lineStatus := 0; lineStatus < n3; lineStatus++ {
		if introvertsCountInner[lineStatus] > introvertsCount || extrovertsCountInner[lineStatus] > extrovertsCount {
			continue
		}
		score := scoreInner[lineStatus] + scoreOuter[lineStatus][lineStatusLast]
		best = max(best, score+dfs(lineStatus, row+1, introvertsCount-introvertsCountInner[lineStatus], extrovertsCount-extrovertsCountInner[lineStatus], m, n3, dp, introvertsCountInner, extrovertsCountInner, scoreInner, scoreOuter))
	}
	dp[lineStatusLast][row][introvertsCount][extrovertsCount] = best
	return best
}

func max(a int, b int) int {
	if a > b {
		return a
	}
	return b
}