Problem
Q-learning learns how good each action is in each state, Q(s, a), from experience alone. After taking action a in state s, receiving reward and landing in state s':
- target = reward, if the episode ended (
done= 1) - target = reward + gamma × max over a' of Q(s', a'), otherwise
- Q(s, a) ← Q(s, a) + alpha × (target - Q(s, a))
Start with every Q value at 0 and apply the transitions in order.
Implement q_learning(num_states, num_actions, transitions, alpha, gamma).
Input
num_states, num_actions: The table size.
transitions: One row [state, action, reward, next_state, done] per step.
alpha: The learning rate.
gamma: The discount factor.
Output
Return the Q table as a matrix with one row per state and one column per action.
Constraints
1 ≤ num_states, num_actions ≤ 100, 0 ≤ number of transitions ≤ 10,000
-100 ≤ reward ≤ 100, done is 0 or 1
0 < alpha ≤ 1, 0 ≤ gamma ≤ 1
Examples
Input:
alpha = 0.5
gamma = 0.9
num_states = 2
num_actions = 2
transitions = [
[0, 1, 0, 1, 0],
[1, 0, 10, 0, 1]
]
Expected output:
expected = [
[0.0, 0.0],
[5.0, 0.0]
]Input:
double alpha = 0.5;
double gamma = 0.9;
long long num_states = 2;
long long num_actions = 2;
std::vector<std::vector<long long>> transitions = {
{0, 1, 0, 1, 0},
{1, 0, 10, 0, 1}
};
Expected output:
std::vector<std::vector<double>> expected = {
{0.0, 0.0},
{5.0, 0.0}
};Input:
double alpha = 0.5;
double gamma = 0.9;
long num_states = 2;
long num_actions = 2;
long[][] transitions = new long[][] {
new long[] {0, 1, 0, 1, 0},
new long[] {1, 0, 10, 0, 1}
};
Expected output:
double[][] expected = new double[][] {
new double[] {0.0, 0.0},
new double[] {5.0, 0.0}
};Explanation
The first step finds no reward yet: target = 0 + 0.9 × max(Q[1]) = 0, so Q[0][1] stays 0. The second step ends the episode with reward 10: Q[1][0] = 0 + 0.5 × (10 - 0) = 5.
Input:
alpha = 0.5
gamma = 0.9
num_states = 2
num_actions = 2
transitions = [
[1, 0, 10, 0, 1],
[0, 1, 0, 1, 0]
]
Expected output:
expected = [
[0.0, 2.25],
[5.0, 0.0]
]Input:
double alpha = 0.5;
double gamma = 0.9;
long long num_states = 2;
long long num_actions = 2;
std::vector<std::vector<long long>> transitions = {
{1, 0, 10, 0, 1},
{0, 1, 0, 1, 0}
};
Expected output:
std::vector<std::vector<double>> expected = {
{0.0, 2.25},
{5.0, 0.0}
};Input:
double alpha = 0.5;
double gamma = 0.9;
long num_states = 2;
long num_actions = 2;
long[][] transitions = new long[][] {
new long[] {1, 0, 10, 0, 1},
new long[] {0, 1, 0, 1, 0}
};
Expected output:
double[][] expected = new double[][] {
new double[] {0.0, 2.25},
new double[] {5.0, 0.0}
};Explanation
Same steps in the other order. Q[1][0] becomes 5 first, so the step from state 0 now sees it: target = 0 + 0.9 × 5 = 4.5 and Q[0][1] = 0.5 × 4.5 = 2.25. The reward has flowed back one state.
How to approach it
Try on your own first, then open one hint at a time.
Hint 1
Each transition updates one table entry using only its current value and the best estimate from the next state.
Hint 2
Keep the Q table initialized to zero, then process transitions in order and update immediately after each one.
Hint 3
If the episode ends, there is no future value to add, and the next-state maximum should not be used.
Expected complexity
O(T * A) time, O(S * A) extra space
References
This challenge is adapted for CodeUtility. Follow the statement and constraints above.