Problem
Value iteration solves a Markov decision process: it finds V(s), the best total discounted reward you can collect starting from state s.
Here the transitions are deterministic. In state s, action a gives reward rewards[s][a] and moves to state next_states[s][a], where -1 means the episode ends. Start with V(s) = 0 for every state, then repeat iterations times, computing every new value from the previous sweep:
V_new(s) = max over a of ( rewards[s][a] + gamma × V(next_states[s][a]) )
where the gamma × V term is 0 when the next state is -1.
Implement value_iteration(next_states, rewards, gamma, iterations).
Input
next_states: For each state, the next state of each of its actions (-1 ends the episode). States may have different numbers of actions.
rewards: The reward of each action, with the same shape as next_states.
gamma: The discount factor.
iterations: The number of sweeps.
Output
Return V after the last sweep, one value per state.
Constraints
1 ≤ states ≤ 200, 1 ≤ actions per state ≤ 10, 0 ≤ iterations ≤ 1,000
-100 ≤ each reward ≤ 100, 0 ≤ gamma ≤ 1
Examples
Input:
gamma = 0.9
rewards = [
[0.0, 0.0],
[0.0, 0.0],
[10.0]
]
iterations = 3
next_states = [
[0, 1],
[0, 2],
[-1]
]
Expected output:
expected = [8.1, 9.0, 10.0]Input:
double gamma = 0.9;
std::vector<std::vector<double>> rewards = {
{0.0, 0.0},
{0.0, 0.0},
{10.0}
};
long long iterations = 3;
std::vector<std::vector<long long>> next_states = {
{0, 1},
{0, 2},
{-1}
};
Expected output:
std::vector<double> expected = {8.1, 9.0, 10.0};Input:
double gamma = 0.9;
double[][] rewards = new double[][] {
new double[] {0.0, 0.0},
new double[] {0.0, 0.0},
new double[] {10.0}
};
long iterations = 3;
long[][] next_states = new long[][] {
new long[] {0, 1},
new long[] {0, 2},
new long[] {-1}
};
Expected output:
double[] expected = new double[] {8.1, 9.0, 10.0};Explanation
Sweep 1: only state 2 can collect a reward, V = [0, 0, 10]. Sweep 2: state 1 moves right to it, 0.9 × 10 = 9, V = [0, 9, 10]. Sweep 3: state 0 gets 0.9 × 9 = 8.1. Result: [8.1, 9, 10].
Input:
gamma = 0.0
rewards = [
[1.0, 2.0],
[3.0, 0.0]
]
iterations = 5
next_states = [
[1, 0],
[0, 1]
]
Expected output:
expected = [2.0, 3.0]Input:
double gamma = 0.0;
std::vector<std::vector<double>> rewards = {
{1.0, 2.0},
{3.0, 0.0}
};
long long iterations = 5;
std::vector<std::vector<long long>> next_states = {
{1, 0},
{0, 1}
};
Expected output:
std::vector<double> expected = {2.0, 3.0};Input:
double gamma = 0.0;
double[][] rewards = new double[][] {
new double[] {1.0, 2.0},
new double[] {3.0, 0.0}
};
long iterations = 5;
long[][] next_states = new long[][] {
new long[] {1, 0},
new long[] {0, 1}
};
Expected output:
double[] expected = new double[] {2.0, 3.0};Explanation
With gamma = 0 the future does not count, so each state's value is just its best immediate reward: [2, 3].
How to approach it
Try on your own first, then open one hint at a time.
Hint 1
Every sweep uses only the values from the previous sweep, not updates made earlier in the same pass.
Hint 2
For each state, compare all actions by combining immediate reward with the discounted value of the next state.
Hint 3
Handle terminal actions carefully: they contribute only their immediate reward when the next state is missing.
Expected complexity
O(iterations * total_actions) time, O(states) extra space
References
This challenge is adapted for CodeUtility. Follow the statement and constraints above.