median_sort.nim

import math
import strutils


proc select_random_pivot_index(anarray: openArray[int], left: int, right: int): int =
  ## ランダムに 軸値 の添字を決定する。
  if left == right:
    return left
  return random(left..right)


proc partition(anarray: var openArray[int], left: int, right: int, pivot_index: int): int =
  ## 最初に軸値 A[pivot_index] を A[right] に移動させる。
  ## 次に軸値 A[right]※ と比較して、下記のように要素を移動させる。
  ## - 小さい場合は A[left, right] の 前半部分
  ## - 大きい場合は A[left, right] の 後半部分
  ## 最後に軸値 A[right] を前半部分の次の要素と入れ替え、軸値がある添字を返す。
  var tmp, store: int

  let pivot = anarray[pivot_index]

  tmp = anarray[right]
  anarray[right] = anarray[pivot_index]
  anarray[pivot_index] = tmp

  store = left
  for i in countup(left, right-1):
    if anarray[i] <= pivot:
      tmp = anarray[i]
      anarray[i] = anarray[store]
      anarray[store] = tmp
      store += 1

  tmp = anarray[right]
  anarray[right] = anarray[store]
  anarray[store] = tmp

  return store


proc select_kth(anarray: var openArray[int], k: int, left: int, right: int): int =
  ## 中央値の添字を求める。
  ## まず軸値となる要素を決める。
  ## 次に軸値を基準とした分割を行い、軸値の添字 p を取得する。
  ## 軸値の添字 p と、 A[left, rigth]において 1<=k<right-left+1となる k 番目に
  ## 小さい値となる k を下記のように比較する。
  ## - p+1 と k が等しい場合、p の位置にある値が中央値である
  ## - p+1 より k が小さい場合、A の k 番目の値は A[left, p] の k 番目の値
  ## - p+1 より k が大きい場合、A の k 番目の値は A[p+1, right] のk-p-1 番目の値
  ##
  ## 軸値の選択が一方に偏るような選択をした場合は O(n**2) となる。
  ## 具体的には既にソート済みの配列に対して、常に一番左側を軸値となるように選
  ## 択する場合が最悪ケースとなる。
  ##
  ## Sell:: https://en.wikipedia.org/wiki/Quickselect
  let idx = select_random_pivot_index(anarray, left, right)
  let pivot_index = partition(anarray, left, right, idx)
  if left+k-1 == pivot_index:
    return pivot_index
  elif left+k-1 < pivot_index:
    return select_kth(anarray, k, left, pivot_index-1)
  else:
    return select_kth(anarray, k-(pivot_index-left+1), pivot_index+1, right)


proc sort*(anarray: var openArray[int], left: int, right: int) =
  if right <= left:
    return
  let mid = (right - left + 1) div 2
  discard select_kth(anarray, mid+1, left, right)
  sort(anarray, left, left+mid-1)
  sort(anarray, left+mid+1, right)


when isMainModule:
  randomize()
  var anarray: array[16, int]
  anarray = [15, 9, 8, 1, 4, 11, 7, 12, 13, 6, 5, 3, 16, 2, 10, 14]
  sort(anarray, 0, 15)
  for v in anarray:
    echo v