# AC
# Runtime: 408 ms, faster than 23.80% of Python3 online submissions for Implement Rand10() Using Rand7().
# Memory Usage: 16.7 MB, less than 90.76% of Python3 online submissions for Implement Rand10() Using Rand7().
class Solution:
    def rand10(self):
        while True:
            a = rand7()
            if a <= 3:
                b = rand7()
                if b <= 5:
                    return b
            elif a <= 6:
                b = rand7()
                if b <= 5:
                    return b + 5

# AC
# Runtime: 376 ms, faster than 54.71% of Python3 online submissions for Implement Rand10() Using Rand7().
# Memory Usage: 16.9 MB, less than 38.54% of Python3 online submissions for Implement Rand10() Using Rand7().
class Solution:
    def rand10(self):
        while True:
            a, b = rand7(),  rand7()
            num = (a - 1) * 7 + b
            if num <= 40:
                return num % 10 + 1

# AC
# Runtime: 344 ms, faster than 92.72% of Python3 online submissions for Implement Rand10() Using Rand7().
# Memory Usage: 16.7 MB, less than 90.76% of Python3 online submissions for Implement Rand10() Using Rand7().
class Solution:
    def rand10(self):
        while True:
            a, b = rand7(),  rand7()
            num = (a - 1) * 7 + b
            if num <= 40: return num % 10 + 1
            a = num - 40
            b = rand7()
            num = (a - 1) * 7 + b
            if num <= 60: return num % 10 + 1
            a = num - 60
            b = rand7()
            num = (a - 1) * 7 + b
            if num <= 20: return num % 10 + 1

# The rand7() API is already defined for you.
rand7_c = 0
rand10_c = 0

def rand7():
    global rand7_c
    rand7_c += 1
    import random
    return random.randint(1, 7)
    
def rand10():
    global rand10_c
    rand10_c += 1
    while True:
        a, b = rand7(), rand7()
        num = (a - 1) * 7 + b
        if num <= 40: return num % 10 + 1
        a = num - 40   # [1, 9]
        b = rand7()
        num = (a - 1) * 7 + b  # [1, 63]
        if num <= 60: return num % 10 + 1
        a = num - 60  # [1, 3]
        b = rand7()
        num = (a - 1) * 7 + b  # [1, 21]
        if num <= 20: return num % 10 + 1

if __name__ == '__main__':
    while True:
        rand10()
        print(f'{rand10_c} {round(rand7_c/rand10_c, 2)}')

import math
from typing import List

class Solution:
    def twoCitySchedCost(self, costs: List[List[int]]) -> int:
        L = range(len(costs))
        from itertools import combinations
        left_set_list = [set(c) for c in combinations(list(L), len(L)//2)]

        min_total = math.inf
        for left_set in left_set_list:
            cost = 0
            for i in L:
                is_left = 1 if i in left_set else 0
                cost += costs[i][is_left]
            min_total = min(min_total, cost)

        return min_total

# AC
# Runtime: 36 ms, faster than 87.77% of Python3 online submissions
# Memory Usage: 14.5 MB, less than 14.84% of Python3 online
from typing import List

class Solution:
    def twoCitySchedCost(self, costs: List[List[int]]) -> int:
        L = range(len(costs))
        cost_diff_lst = [(i, costs[i][0] - costs[i][1]) for i in L]
        cost_diff_lst.sort(key=lambda x: x[1])

        total_cost = 0
        for c, (idx, _) in enumerate(cost_diff_lst):
            is_left = 0 if c < len(L) // 2 else 1
            total_cost += costs[idx][is_left]

        return total_cost

from ortools.sat.python import cp_model

costs = [[515,563],[451,713],[537,709],[343,819],[855,779],[457,60],[650,359],[631,42]]

I = range(len(costs))

model = cp_model.CpModel()
x = []
total_cost = model.NewIntVar(0, 10000, 'total_cost')
for i in I:
    t = []
    for j in range(2):
        t.append(model.NewBoolVar('x[%i,%i]' % (i, j)))
    x.append(t)

# Constraints
# Each person must be assigned to at exact one city
[model.Add(sum(x[i][j] for j in range(2)) == 1) for i in I]
# equal number of person assigned to two cities
model.Add(sum(x[i][0] for i in I) == (len(I) // 2))

# Total cost
model.Add(total_cost == sum(x[i][0] * costs[i][0] + x[i][1] * costs[i][1] for i in I))
model.Minimize(total_cost)

solver = cp_model.CpSolver()
status = solver.Solve(model)

if status == cp_model.OPTIMAL:
    print('Total min cost = %i' % solver.ObjectiveValue())
    print()
    for i in I:
        for j in range(2):
            if solver.Value(x[i][j]) == 1:
                print('People ', i, ' assigned to city ', j, '  Cost = ', costs[i][j])

import pulp

costs = [[259,770],[448,54],[926,667],[184,139],[840,118],[577,469]]   # 1859


I = range(len(costs))

items=[i for i in I]
city_a = pulp.LpVariable.dicts('left', items, 0, 1, pulp.LpBinary)
city_b = pulp.LpVariable.dicts('right', items, 0, 1, pulp.LpBinary)

m = pulp.LpProblem("Two Cities", pulp.LpMinimize)

m += pulp.lpSum((costs[i][0] * city_a[i] + costs[i][1] * city_b[i]) for i in items)

# Constraints
# Each person must be assigned to at exact one city
for i in I:
    m += pulp.lpSum([city_a[i] + city_b[i]]) == 1
# create a binary variable to state that a table setting is used
m += pulp.lpSum(city_a[i] for i in I) == (len(I) // 2)

m.solve()

total = 0
for i in I:
    if city_a[i].value() == 1.0:
        total += costs[i][0]
    else:
        total += costs[i][1]

print("Total cost {}".format(total))

def search_brute_force(initial: List, transition: List, L: int) -> Tuple[float, Tuple]:
    from itertools import combinations_with_replacement
    v = [0, 1, 2]
    path_all = combinations_with_replacement(v, L)

    max_prop = 0.0
    max_route = None
    prob = 0.0
    for path in list(path_all):
        for idx, v in enumerate(path):
            if idx == 0:
                prob = initial[v]  # reset to initial state
            else:
                prev_v = path[idx-1]
                prob *= transition[prev_v][v]
        if prob > max_prop:
            max_prop = max(max_prop, prob)
            max_route = path
    return max_prop, max_route

def search_greedy(initial: List, transition: List, L: int) -> Tuple[float, Tuple]:
    import numpy as np
    max_route = []
    max_prop = 0.0
    states = np.array(initial)

    prev_max_v = None
    for l in range(0, L):
        max_v = np.argmax(states)
        max_route.append(max_v)
        if l == 0:
            max_prop = initial[max_v]
        else:
            max_prop = max_prop * transition[prev_max_v][max_v]
        states = max_prop * states
        prev_max_v = max_v

    return max_prop, max_route

def search_beam(initial: List, transition: List, L: int, K: int) -> Tuple[float, Tuple]:
    N = len(initial)
    from queue import PriorityQueue
    current_q = PriorityQueue()
    next_q = PriorityQueue()

    from functools import total_ordering
    @total_ordering
    class PQItem(object):
        def __init__(self, prob, route):
            self.prob = prob
            self.route = route
            self.last_v = int(route[-1])

        def __eq__(self, other):
            return self.prob == other.prob

        def __lt__(self, other):
            return self.prob > other.prob

    for v in range(N):
        next_q.put(PQItem(initial[v], str(v)))

    for l in range(1, L):
        current_q = next_q
        next_q = PriorityQueue()
        k = K
        while not current_q.empty() and k > 0:
            item = current_q.get()
            prob, route, prev_v = item.prob, item.route, item.last_v
            k -= 1
            for v in range(N):
                nextItem = PQItem(prob * transition[prev_v][v], route + str(v))
                next_q.put(nextItem)

    max_item = next_q.get()

    return max_item.prob, list(map(lambda x: int(x), max_item.route))

def search_dp(initial: List, transition: List, L: int) -> float:
    N = len(initial)
    dp = [[0.0 for c in range(N)] for r in range(L)]
    dp[0] = initial[:]

    for l in range(1, L):
        for v in range(N):
            for prev_v in range(N):
                dp[l][v] = max(dp[l][v], dp[l - 1][prev_v] * transition[prev_v][v])

    return max(dp[L-1])

def search_prob_greedy(initial: List, transition: List, L: int) -> Tuple[float, Tuple]:
    import random
    N = len(initial)
    max_route = []
    max_prop = 0.0
    vertices = [i for i in range(N)]
    prob = initial[:]

    for l in range(0, L):
        v_lst = random.choices(vertices, prob)
        v = v_lst[0]
        max_route.append(v)
        max_prop = prob[v]
        prob = [prob[v] * transition[v][v_target] for v_target in range(N)]

    return max_prop, max_route

from torch.utils.data import Dataset

class TSPDataset(Dataset):
	"each data item of form (input, input_len, output_in, output_out, output_len)"
	data: List[Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]]
	
	def __len__(self):
		return len(self.data)

	def __getitem__(self, index):
		input, input_len, output_in, output_out, output_len = self.data[index]
		return input, input_len, output_in, output_out, output_len

class RNNEncoder(nn.Module):
	rnn: Union[nn.LSTM, nn.GRU, nn.RNN]

	def __init__(self, rnn_type: str, bidirectional: bool, num_layers: int, input_size: int, hidden_size: int, dropout: float):
		super(RNNEncoder, self).__init__()
		if bidirectional:
			assert hidden_size % 2 == 0
			hidden_size = hidden_size // 2
		self.rnn = rnn_init(rnn_type, input_size=input_size, hidden_size=hidden_size, bidirectional=bidirectional,num_layers=num_layers, dropout=dropout)

	def forward(self, src: Tensor, src_lengths: Tensor, hidden: Tensor = None) -> Tuple[Tensor, Tensor]:
		lengths = src_lengths.view(-1).tolist()
		packed_src = pack_padded_sequence(src, lengths)
		memory_bank, hidden_final = self.rnn(packed_src, hidden)
		memory_bank = pad_packed_sequence(memory_bank)[0]
		return memory_bank, hidden_final

class Attention(nn.Module):
	linear_out: nn.Linear

	def __init__(self, dim: int):
		super(Attention, self).__init__()
		self.linear_out = nn.Linear(dim * 2, dim, bias=False)

	def score(self, src: Tensor, target: Tensor) -> Tensor:
		batch_size, src_len, dim = src.size()
		_, target_len, _ = target.size()
		target_ = target
		src_ = src.transpose(1, 2)
		return torch.bmm(target_, src_)

	def forward(self, src: Tensor, target: Tensor, src_lengths: Tensor) -> Tuple[Tensor, Tensor]:
		assert target.dim() == 3

		batch_size, src_len, dim = src.size()
		_, target_len, _ = target.size()

		align_score = self.score(src, target)

		mask = sequence_mask(src_lengths)
		# (batch_size, max_len) -> (batch_size, 1, max_len)
		mask = mask.unsqueeze(1)
		align_score.data.masked_fill_(~mask, -float('inf'))
		align_score = F.softmax(align_score, -1)

		c = torch.bmm(align_score, src)

		concat_c = torch.cat([c, target], -1)
		attn_h = self.linear_out(concat_c)

		return attn_h, align_score

class TSPSolver:
    g: Graph

    def __init__(self, g: Graph):
        self.g = g

    def solve(self) -> int:
        """
        :param v:
        :param state:
        :return: -1 means INF
        """
        N = self.g.v_num
        dp = [[INT_INF for c in range(N)] for r in range(1 << N)]

        dp[(1 << N) - 1][0] = 0

        for state in range((1 << N) - 2, -1, -1):
            for v in range(N):
                for u in range(N):
                    if ((state >> u) & 1) == 0:
                        if dp[state | 1 << u][u] != INT_INF and self.g.edges[v][u] != INT_INF:
                            if dp[state][v] == INT_INF:
                                dp[state][v] = dp[state | 1 << u][u] + self.g.edges[v][u]
                            else:
                                dp[state][v] = min(dp[state][v], dp[state | 1 << u][u] + self.g.edges[v][u])
        return dp[0][0]

import matplotlib.pyplot as plt
points='0.607122 0.664447 0.953593 0.021519 0.757626 0.921024 0.586376 0.433565 0.786837 0.052959 0.016088 0.581436 0.496714 0.633571 0.227777 0.971433 0.665490 0.074331 0.383556 0.104392'
float_list = list(map(lambda x: float(x), points.split(' ')))

x,y = [],[]
for idx, p in enumerate(float_list):
  if idx % 2 == 0:
    x.append(p)
  else:
    y.append(p)

for i in range(0, len(x)):
  for j in range(0, len(x)):
    if i == j:
      continue
    plt.plot((x[i],x[j]),(y[i],y[j]))

plt.show()

tour_str = '1 3 8 6 10 9 5 2 4 7 1'
tour = list(map(lambda x: int(x), tour_str.split(' ')))

for i in range(0, len(tour)-1):
  p1 = tour[i] - 1
  p2 = tour[i + 1] - 1
  plt.plot((x[p1],x[p2]),(y[p1],y[p2]))
plt.show()

g: Graph = Graph(N)
for v in range(N):
	for u in range(N):
		diff_x = coordinates[v][0] - coordinates[u][0]
		diff_y = coordinates[v][1] - coordinates[u][1]
		dist: float = math.sqrt(diff_x * diff_x + diff_y * diff_y)
		g.setDist(u, v, dist)
		g.setDist(v, u, dist)

ret: float = FLOAT_INF
u_min: int = -1
	for u in range(self.g.v_num):
		if (state & (1 << u)) == 0:
			s: float = self._recurse(u, state | 1 << u)
				if s + edges[v][u] < ret:
					ret = s + edges[v][u]
					u_min = u
	dp[state][v] = ret
	self.parent[state][v] = u_min

def _form_tour(self):
	self.tour = [0]
	bit = 0
	v = 0
	for _ in range(self.g.v_num - 1):
		v = self.parent[bit][v]
		self.tour.append(v)
		bit = bit | (1 << v)
	self.tour.append(0)

tsp: TSPSolver = TSPSolver(g)
tsp.solve()

output_dist: float = 0.0
output_tour = list(map(lambda x: int(x) - 1, output.split(' ')))
for v in range(1, len(output_tour)):
	pre_v = output_tour[v-1]
	curr_v = output_tour[v]
	diff_x = coordinates[pre_v][0] - coordinates[curr_v][0]
	diff_y = coordinates[pre_v][1] - coordinates[curr_v][1]
	dist: float = math.sqrt(diff_x * diff_x + diff_y * diff_y)
	output_dist += dist

	passed = abs(tsp.dist - output_dist) < 10e-5
	if passed:
		print(f'passed dist={tsp.tour}')
	else:
		print(f'Min Tour Distance = {output_dist}, Computed Tour Distance = {tsp.dist}, Expected Tour = {output_tour}, Result = {tsp.tour}')

# AC
# Runtime: 36 ms, faster than 91.78% of Python3 online submissions for Permutations.
# Memory Usage: 13.9 MB, less than 66.52% of Python3 online submissions for Permutations.

from itertools import permutations
from typing import List


class Solution:
    def permute(self, nums: List[int]) -> List[List[int]]:
        return [p for p in permutations(nums)]

# AC
# Runtime: 84 ms, faster than 95.43% of Python3 online submissions for Combinations.
# Memory Usage: 15.2 MB, less than 68.98% of Python3 online submissions for Combinations.
from itertools import combinations
from typing import List

class Solution:
    def combine(self, n: int, k: int) -> List[List[int]]:
        return [c for c in combinations(list(range(1, n + 1)), k)]

[lst for lst in product([1, 2, 3], ['a', 'b'])]
[(i, s) for i in [1, 2, 3] for s in ['a', 'b']]

# AC
# Runtime: 24 ms, faster than 94.50% of Python3 online submissions for Letter Combinations of a Phone Number.
# Memory Usage: 13.7 MB, less than 83.64% of Python3 online submissions for Letter Combinations of a Phone Number.

from itertools import product
from typing import List


class Solution:
    def letterCombinations(self, digits: str) -> List[str]:
        if digits == "":
            return []
        mapping = {'2':'abc', '3':'def', '4':'ghi', '5':'jkl', '6':'mno', '7':'pqrs', '8':'tuv', '9':'wxyz'}
        iter_dims = [mapping[i] for i in digits]

        result = []
        for lst in product(*iter_dims):
            result.append(''.join(lst))

        return result

# AC
# Runtime: 28 ms, faster than 85.56% of Python3 online submissions for Fibonacci Number.
# Memory Usage: 13.8 MB, less than 58.41% of Python3 online submissions for Fibonacci Number.

class Solution:
    def fib(self, N: int) -> int:
        if N <= 1:
            return N
        i = 2
        for fib in self.fib_next():
            if i == N:
                return fib
            i += 1
            
    def fib_next(self):
        f_last2, f_last = 0, 1
        while True:
            f = f_last2 + f_last
            f_last2, f_last = f_last, f
            yield f

def outer_generator():
    for i in inner_generator():
        yield i

# AC
# Runtime: 48 ms, faster than 90.31% of Python3 online submissions for Kth Smallest Element in a BST.
# Memory Usage: 17.9 MB, less than 14.91% of Python3 online submissions for Kth Smallest Element in a BST.

class Solution:
    def kthSmallest(self, root: TreeNode, k: int) -> int:
        def ordered_iter(node):
            if node:
                for sub_node in ordered_iter(node.left):
                    yield sub_node
                yield node
                for sub_node in ordered_iter(node.right):
                    yield sub_node

        for node in ordered_iter(root):
            k -= 1
            if k == 0:
                return node.val

# AC
# Runtime: 56 ms, faster than 63.74% of Python3 online submissions for Kth Smallest Element in a BST.
# Memory Usage: 17.7 MB, less than 73.33% of Python3 online submissions for Kth Smallest Element in a BST.

class Solution:
    def kthSmallest(self, root: TreeNode, k: int) -> int:
        def ordered_iter(node):
            if node:
                yield from ordered_iter(node.left)
                yield node
                yield from ordered_iter(node.right)

        for node in ordered_iter(root):
            k -= 1
            if k == 0:
                return node.val

# AC
# Runtime: 112 ms, faster than 57.59% of Python3 online submissions for 24 Game.
# Memory Usage: 13.7 MB, less than 85.60% of Python3 online submissions for 24 Game.

import math
from itertools import permutations, product
from typing import List

class Solution:

    def iter_trees(self, op1, op2, op3, a, b, c, d):
        yield op1(op2(a, b), op3(c, d))
        yield op1(a, op2(op3(b, c), d))
        yield op1(a, op2(b, op3(c, d)))
        yield op1(op2(a, op3(b, c)), d)

    def judgePoint24(self, nums: List[int]) -> bool:
        mul = lambda x, y: x * y
        plus = lambda x, y: x + y
        div = lambda x, y: x / y if y != 0 else math.inf
        minus = lambda x, y: x - y

        op_lst = [plus, minus, mul, div]

        for ops in product(op_lst, repeat=3):
            for val in permutations(nums):
                for v in self.iter_trees(ops[0], ops[1], ops[2], val[0], val[1], val[2], val[3]):
                    if abs(v - 24) < 0.0001:
                        return True
        return False

# AC
# Runtime: 116 ms, faster than 56.23% of Python3 online submissions for 24 Game.
# Memory Usage: 13.9 MB, less than 44.89% of Python3 online submissions for 24 Game.

import math
from itertools import combinations, product, permutations
from typing import List

class Solution:

    def judgePoint24(self, nums: List[int]) -> bool:
        mul = lambda x, y: x * y
        plus = lambda x, y: x + y
        div = lambda x, y: x / y if y != 0 else math.inf
        minus = lambda x, y: x - y

        op_lst = [plus, minus, mul, div]

        def recurse(lst: List[int]):
            if len(lst) == 2:
                for op, values in product(op_lst, permutations(lst)):
                    yield op(values[0], values[1])
            else:
                # choose 2 indices from lst of length n
                for choosen_idx_lst in combinations(list(range(len(lst))), 2):
                    # remaining indices not choosen (of length n-2)
                    idx_remaining_set = set(list(range(len(lst)))) - set(choosen_idx_lst)

                    # remaining values not choosen (of length n-2)
                    value_remaining_lst = list(map(lambda x: lst[x], idx_remaining_set))
                    for op, idx_lst in product(op_lst, permutations(choosen_idx_lst)):
                        # 2 choosen values are lst[idx_lst[0]], lst[idx_lst[1]
                        value_remaining_lst.append(op(lst[idx_lst[0]], lst[idx_lst[1]]))
                        yield from recurse(value_remaining_lst)
                        value_remaining_lst = value_remaining_lst[:-1]

        for v in recurse(nums):
            if abs(v - 24) < 0.0001:
                return True

from itertools import permutations
v = [1,2,3]
p = permutations(v)
for t in list(p):
  print([0] + list(t) + [0])

[0, 1, 2, 3, 0]
[0, 1, 3, 2, 0]
[0, 2, 1, 3, 0]
[0, 2, 3, 1, 0]
[0, 3, 1, 2, 0]
[0, 3, 2, 1, 0]

INT_INF = -1

if s != INT_INF and edges[v][u] != INT_INF:
    if ret == INT_INF:
        ret = s + edges[v][u]
    else:
        ret = min(ret, s + edges[v][u])

from typing import List

INT_INF = -1

class Graph:
    v_num: int
    edges: List[List[int]]

    def __init__(self, v_num: int):
        self.v_num = v_num
        self.edges = [[INT_INF for c in range(v_num)] for r in range(v_num)]
    
    def setDist(self, src: int, dest: int, dist: int):
        self.edges[src][dest] = dist


class TSPSolver:
    g: Graph
    dp: List[List[int]]

    def __init__(self, g: Graph):
        self.g = g
        self.dp = [[None for c in range(g.v_num)] for r in range(1 << g.v_num)]
    
    def solve(self) -> int:
        return self._recurse(0, 0)
    
    def _recurse(self, v: int, state: int) -> int:
        """
    
        :param v:
        :param state:
        :return: -1 means INF
        """
        dp = self.dp
        edges = self.g.edges
    
        if dp[state][v] is not None:
            return dp[state][v]
    
        if (state == (1 << self.g.v_num) - 1) and (v == 0):
            dp[state][v] = 0
            return dp[state][v]
    
        ret: int = INT_INF
        for u in range(self.g.v_num):
            if (state & (1 << u)) == 0:
                s: int = self._recurse(u, state | 1 << u)
                if s != INT_INF and edges[v][u] != INT_INF:
                    if ret == INT_INF:
                        ret = s + edges[v][u]
                    else:
                        ret = min(ret, s + edges[v][u])
        dp[state][v] = ret
        return ret


def main():
    V, E = map(int, input().split())
    g: Graph = Graph(V)
    for _ in range(E):
        src, dest, dist = map(int, input().split())
        g.setDist(src, dest, dist)

    tsp: TSPSolver = TSPSolver(g)
    print(tsp.solve())


if __name__ == "__main__":
    main()

def simulate_bruteforce(n: int) -> bool:
    """
    Simulates one round. Unbounded time complexity.
    :param n: total number of seats
    :return: True if last one has last seat, otherwise False
    """

    seats = [False for _ in range(n)]

    for i in range(n-1):
        if i == 0:  # first one, always random
            seats[random.randint(0, n - 1)] = True
        else:
            if not seats[i]:  # i-th has his seat
                seats[i] = True
            else:
                while True:
                    rnd = random.randint(0, n - 1) # random until no conflicts
                    if not seats[rnd]:
                        seats[rnd] = True
                        break
    return not seats[n-1]

class Solution:
    def nthPersonGetsNthSeat(self, n: int) -> float:
        return 1.0 if n == 1 else 0.5

def simulate_online(n: int) -> bool:
    """
    Simulates one round of complexity O(N).
    :param n: total number of seats
    :return: True if last one has last seat, otherwise False
    """

    seats = [i for i in range(n)]

    def swap(i, j):
        tmp = seats[i]
        seats[i] = seats[j]
        seats[j] = tmp

    # for each person, the seats array idx available are [i, n-1]
    for i in range(n-1):
        if i == 0:  # first one, always random
            rnd = random.randint(0, n - 1)
            swap(rnd, 0)
        else:
            if seats[i] == i:  # i-th still has his seat
                pass
            else:
                rnd = random.randint(i, n - 1)  # selects idx from [i, n-1]
                swap(rnd, i)
    return seats[n-1] == n - 1

def simulate(n: int) -> bool:
    """
    Simulates one round of complexity O(N).
    :param n: total number of seats
    :return: True if last one has last seat, otherwise False
    """

    seats = [False for _ in range(n)]

    for i in range(n-1):
        if i == 0:  # first one, always random
            rnd = random.randint(0, n - 1)
            seats[rnd] = True
        else:
            if not seats[i]:  # i-th still has his seat
                seats[i] = True
            else:
                # 0 must not be available, now we have 0 and [i+1, n-1],
                rnd = random.randint(i, n - 1)
                if rnd == i:
                    seats[0] = True
                else:
                    seats[rnd] = True
    return not seats[n-1]

class TreeNode:
    """
    MCTS Tree Node
    """

    c_puct: ClassVar[int] = 5  # class-wise global param c_puct, exploration weight factor.

    _parent: TreeNode
    _children: Dict[int, TreeNode]  # map from action to TreeNode
    _visit_num: int
    _Q: float   # Q value of the node, which is the mean action value.
    _prior: float

class TreeNode:

    def get_puct(self) -> float:
        """
        Computes AlphaGo Zero PUCT (polynomial upper confidence trees) of the node.

        :return: Node PUCT value.
        """
        U = (TreeNode.c_puct * self._prior * np.sqrt(self._parent._visit_num) / (1 + self._visit_num))
        return self._Q + U

class TreeNode:

    def select(self) -> Tuple[Pos, TreeNode]:
        """
        Selects an action(Pos) having max UCB value.

        :return: Action and corresponding node
        """
        return max(self._children.items(), key=lambda act_node: act_node[1].get_puct())

class TreeNode:

    def propagate_to_root(self, leaf_value: float):
        """
        Updates current node with observed leaf_value and propagates to root node.

        :param leaf_value:
        :return:
        """
        if self._parent:
            self._parent.propagate_to_root(-leaf_value)
        self._update(leaf_value)

    def _update(self, leaf_value: float):
        """
        Updates the node by newly observed leaf_value.

        :param leaf_value:
        :return:
        """
        self._visit_num += 1
        # new Q is updated towards deviation from existing Q
        self._Q += 0.5 * (leaf_value - self._Q)

class MCTSAlphaGoZeroPlayer(BaseAgent):

    def _next_step_play_act_probs(self, game: ConnectNGame) -> Tuple[List[Pos], ActionProbs]:
        """
        For the given game status, run playouts number of times specified by self._playout_num.
        Returns the action distribution according to AlphaGo Zero MCTS play formula.

        :param game:
        :return: actions and their probability
        """

        for n in range(self._playout_num):
            self._playout(copy.deepcopy(game))

        act_visits = [(act, node._visit_num) for act, node in self._current_root._children.items()]
        acts, visits = zip(*act_visits)
        act_probs = softmax(1.0 / MCTSAlphaGoZeroPlayer.temperature * np.log(np.array(visits) + 1e-10))

        return acts, act_probs

class MCTSAlphaGoZeroPlayer(BaseAgent):

    def get_action(self, board: PyGameBoard) -> Pos:
        """
        Method defined in BaseAgent.

        :param board:
        :return: next move for the given game board.
        """
        return self._get_action(copy.deepcopy(board.connect_n_game))[0]

    def _get_action(self, game: ConnectNGame) -> Tuple[MoveWithProb]:
        epsilon = 0.25
        avail_pos = game.get_avail_pos()
        move_probs: ActionProbs = np.zeros(game.board_size * game.board_size)
        assert len(avail_pos) > 0

        # the pi defined in AlphaGo Zero paper
        acts, act_probs = self._next_step_play_act_probs(game)
        move_probs[list(acts)] = act_probs
        if self._is_training:
            # add Dirichlet Noise when training in favour of exploration
            p_ = (1-epsilon) * act_probs + epsilon * np.random.dirichlet(0.3 * np.ones(len(act_probs)))
            move = np.random.choice(acts, p=p_)
            assert move in game.get_avail_pos()
        else:
            move = np.random.choice(acts, p=act_probs)

        self.reset()
        return move, move_probs

class MCTSAlphaGoZeroPlayer(BaseAgent):

    def self_play_one_game(self, game: ConnectNGame) \
            -> List[Tuple[NetGameState, ActionProbs, NDArray[(Any), np.float]]]:
        """

        :param game:
        :return:
            Sequence of (s, pi, z) of a complete game play. The number of list is the game play length.
        """

        states: List[NetGameState] = []
        probs: List[ActionProbs] = []
        current_players: List[np.float] = []

        while not game.game_over:
            move, move_probs = self._get_action(game)
            states.append(convert_game_state(game))
            probs.append(move_probs)
            current_players.append(game.current_player)
            game.move(move)

        current_player_z = np.zeros(len(current_players))
        current_player_z[np.array(current_players) == game.game_result] = 1.0
        current_player_z[np.array(current_players) == -game.game_result] = -1.0
        self.reset()

        return list(zip(states, probs, current_player_z))

class MCTSAlphaGoZeroPlayer(BaseAgent):
    
    def _playout(self, game: ConnectNGame):
        """
        From current game status, run a sequence down to a leaf node, either because game ends or unexplored node.
        Get the leaf value of the leaf node, either the actual reward of game or action value returned by policy net.
        And propagate upwards to root node.

        :param game:
        """
        player_id = game.current_player

        node = self._current_root
        while True:
            if node.is_leaf():
                break
            act, node = node.select()
            game.move(act)

        # now game state is a leaf node in the tree, either a terminal node or an unexplored node
        act_and_probs: Iterator[MoveWithProb]
        act_and_probs, leaf_value = self._policy_value_net.policy_value_fn(game)

        if not game.game_over:
            # case where encountering an unexplored leaf node, update leaf_value estimated by policy net to root
            for act, prob in act_and_probs:
                game.move(act)
                child_node = node.expand(act, prob)
                game.undo()
        else:
            # case where game ends, update actual leaf_value to root
            if game.game_result == ConnectNGame.RESULT_TIE:
                leaf_value = ConnectNGame.RESULT_TIE
            else:
                leaf_value = 1 if game.game_result == player_id else -1
            leaf_value = float(leaf_value)

        # Update leaf_value and propagate up to root node
        node.propagate_to_root(-leaf_value)

NetGameState = NDArray[(4, Any, Any), np.int]


def convert_game_state(game: ConnectNGame) -> NetGameState:
    """
    Converts game state to type NetGameState as ndarray.

    :param game:
    :return:
        Of shape 4 * board_size * board_size.
        [0] is current player positions.
        [1] is opponent positions.
        [2] is last move location.
        [3] all 1 meaning move by black player, all 0 meaning move by white.
    """
    state_matrix = np.zeros((4, game.board_size, game.board_size))

    if game.action_stack:
        actions = np.array(game.action_stack)
        move_curr = actions[::2]
        move_oppo = actions[1::2]
        for move in move_curr:
            state_matrix[0][move] = 1.0
        for move in move_oppo:
            state_matrix[1][move] = 1.0
        # indicate the last move location
        state_matrix[2][actions[-1]] = 1.0
    if len(game.action_stack) % 2 == 0:
        state_matrix[3][:, :] = 1.0  # indicate the colour to play
    return state_matrix[:, ::-1, :]

def backward_step(self, state_batch: List[NetGameState], probs_batch: List[ActionProbs],
               value_batch: List[NDArray[(Any), np.float]], lr) -> Tuple[float, float]:
    if self.use_gpu:
        state_batch = Variable(torch.FloatTensor(state_batch).cuda())
        probs_batch = Variable(torch.FloatTensor(probs_batch).cuda())
        value_batch = Variable(torch.FloatTensor(value_batch).cuda())
    else:
        state_batch = Variable(torch.FloatTensor(state_batch))
        probs_batch = Variable(torch.FloatTensor(probs_batch))
        value_batch = Variable(torch.FloatTensor(value_batch))

    self.optimizer.zero_grad()
    for param_group in self.optimizer.param_groups:
        param_group['lr'] = lr

    log_act_probs, value = self.policy_value_net(state_batch)
    # loss = (z - v)^2 - pi*T * log(p) + c||theta||^2
    value_loss = F.mse_loss(value.view(-1), value_batch)
    policy_loss = -torch.mean(torch.sum(probs_batch * log_act_probs, 1))
    loss = value_loss + policy_loss
    loss.backward()
    self.optimizer.step()
    entropy = -torch.mean(torch.sum(torch.exp(log_act_probs) * log_act_probs, 1))
    return loss.item(), entropy.item()