fork(1) download
copy 
  1. # your code goes here
  2. def log2(n):
  3.   i = -1
  4.   while(n):
  5.     i += 1
  6.     n >>= 1
  7.   return i
  8.  
  9. # https://c...content-available-to-author-only...s.com/string/suffix-array.html
  10. def sort_cyclic_shifts(s):
  11.   n = len(s)
  12.   alphabet = 256
  13.   cs = []
  14.  
  15.   p = [0] * n
  16.   c = [0] * n
  17.   cnt = [0] * max(alphabet, n + 1)
  18.  
  19.   for i in range(n):
  20.     cnt[ord(s[i])] += 1
  21.   for i in range(1, alphabet):
  22.     cnt[i] += cnt[i-1]
  23.   for i in range(n):
  24.     cnt[ord(s[i])] -= 1
  25.     p[cnt[ord(s[i])]] = i
  26.   c[p[0]] = 0
  27.   classes = 1
  28.   for i in range(1, n):
  29.     if s[p[i]] != s[p[i-1]]:
  30.       classes += 1
  31.     c[p[i]] = classes - 1
  32.  
  33.   cs.append(c[:])
  34.  
  35.   pn = [0] * n
  36.   cn = [0] * n
  37.   h = 0
  38.  
  39.   while (1 << h) < n:
  40.     for i in range(n):
  41.       pn[i] = p[i] - (1 << h)
  42.       if pn[i] < 0:
  43.         pn[i] += n
  44.  
  45.     for i in range(0, classes):
  46.       cnt[i] = 0
  47.  
  48.     for i in range(n):
  49.       cnt[c[pn[i]]] += 1
  50.     for i in range(1, classes):
  51.       cnt[i] += cnt[i-1]
  52.     for i in range(n-1, -1, -1):
  53.       cnt[c[pn[i]]] -= 1
  54.       p[cnt[c[pn[i]]]] = pn[i]
  55.     cn[p[0]] = 0
  56.     classes = 1
  57.     for i in range(i, n):
  58.       cur = c[p[i]], c[(p[i] + (1 << h)) % n]
  59.       prev = c[p[i-1]], c[(p[i-1] + (1 << h)) % n]
  60.       if cur != prev:
  61.         classes += 1
  62.       cn[p[i]] = classes - 1
  63.     c = cn
  64.     cs.append(c[:])
  65.     h += 1
  66.  
  67.   return p, cs
  68.  
  69. # https://c...content-available-to-author-only...s.com/string/suffix-array.html
  70. def suffix_array_construction(s):
  71.   s += "$"
  72.   sorted_shifts, cs = sort_cyclic_shifts(s)
  73.   return sorted_shifts[1:], cs
  74.  
  75. # https://c...content-available-to-author-only...s.com/string/suffix-array.html
  76. def compare(i, j, l, k, n, c):
  77.   a = c[k][i], c[k][(i+l-(1 << k))%n]
  78.   b = c[k][j], c[k][(j+l-(1 << k))%n]
  79.   if a == b:
  80.     return 0
  81.   elif a < b:
  82.     return -1
  83.   return 1
  84.  
  85.  
  86. ## MAIN FUNCTION
  87.  
  88.  
  89. def f(s, k):
  90.   debug = 0
  91.  
  92.   n = len(s)
  93.  
  94.   # Best prefix
  95.   best_first = s[k]
  96.   best_second = s[k+1]
  97.   first_idxs = [k]
  98.   second_idxs = [k + 1]
  99.  
  100.   for i in range(k - 1, -1, -1):
  101.     if s[i] <= best_first:
  102.       best_first = s[i]
  103.       # We only need one leftmost index
  104.       first_idxs = [i]
  105.   for i in range(k, first_idxs[0], -1):
  106.     if (s[i] < best_second):
  107.       best_second = s[i]
  108.       second_idxs = [i]
  109.     elif s[i] == best_second:
  110.       second_idxs.append(i)
  111.  
  112.   second_idxs = list(reversed(second_idxs))
  113.  
  114.   # Best suffix
  115.   # For each of l deletions,
  116.   # we can place the last
  117.   # character anywhere ahead
  118.   # of the penultimate.
  119.   last_idxs = {(n - 2): [n - 1]}
  120.   best_last = s[n - 1]
  121.   for l in range(2, k + 2):
  122.     idx = n - l
  123.     if s[idx] < best_last:
  124.       best_last = s[idx]
  125.       last_idxs[n - 1 - l] = [idx]
  126.     else:
  127.       last_idxs[n - 1 - l] = last_idxs[n - l]
  128.  
  129.   p, cs = suffix_array_construction(s)
  130.  
  131.   second_idx = 0
  132.  
  133.   if debug:
  134.     print(first_idxs, second_idxs, last_idxs)
  135.  
  136.   while first_idxs[0] >= second_idxs[second_idx]:
  137.     second_idx += 1
  138.  
  139.   prefix_end = second_idxs[second_idx]
  140.   num_deleted = prefix_end - 1
  141.   remaining = k - num_deleted
  142.   suffix_start = n - remaining - 2
  143.   best = (prefix_end + 1, suffix_start - 1)
  144.  
  145.   while second_idx < len(second_idxs):
  146.     prefix_end = second_idxs[second_idx]
  147.     num_deleted = prefix_end - 1
  148.     remaining = k - num_deleted
  149.     suffix_start = n - remaining - 2
  150.  
  151.     len_candidate_middle = suffix_start - 1 - prefix_end
  152.  
  153.     # The prefixes are all equal.
  154.     # We need to compare the middle
  155.     # and suffix.
  156.     # compare(i, j, l, k, n, c)
  157.     len_best_middle = best[1] - best[0] + 1
  158.     l = min(len_candidate_middle, len_best_middle)
  159.  
  160.     # Compare middles
  161.     comp = compare(best[0], prefix_end + 1, l, log2(l), n + 1, cs)
  162.  
  163.     # Candidate is better
  164.     if comp == 1:
  165.       best = (prefix_end + 1, suffix_start - 1)
  166.     elif comp == 0:
  167.       # Compare suffix of candidate with
  168.       # substring at the comparable position
  169.       # of best.
  170.       [last_idx] = last_idxs[suffix_start]
  171.       candidate_suffix = s[suffix_start] + s[last_idx]
  172.  
  173.       if len_candidate_middle < len_best_middle:
  174.         # One character of best's suffix
  175.         if len_candidate_middle + 1 == len_best_middle:
  176.           to_compare = s[best[1]] + s[best[1] + 1]
  177.         # None of best's suffix
  178.         else:
  179.           idx = best[0] + len_candidate_middle
  180.           to_compare = s[idx] + s[idx + 1]
  181.         # If the candidate suffix is equal
  182.         # to best's equivalent, the candidate
  183.         # wins since it's shorter.
  184.         if candidate_suffix <= to_compare:
  185.           best = (prefix_end + 1, suffix_start - 1)
  186.  
  187.       elif len_candidate_middle == len_best_middle:
  188.         idx = best[1] + 1
  189.         to_compare = s[idx] + s[last_idxs[idx][0]]
  190.         if candidate_suffix < to_compare:
  191.           best = (prefix_end + 1, suffix_start - 1)
  192.  
  193.       # len_best_middle < len_candidate_middle 
  194.       else:
  195.         # One character of candidate's suffix
  196.         if len_best_middle + 1 == len_candidate_middle:
  197.           to_compare = s[suffix_start - 1] + s[suffix_start]
  198.         # None of candidates's suffix
  199.         else:
  200.           idx = prefix_end + 1 + len_best_middle
  201.           to_compare = s[idx] + s[idx + 1]
  202.  
  203.         if candidate_suffix < to_compare:
  204.           best = (prefix_end + 1, suffix_start - 1)
  205.  
  206.     second_idx += 1
  207.  
  208.   prefix = s[first_idxs[0]] + s[second_idxs[second_idx-1]]
  209.   middle = s[best[0]:best[1] + 1]
  210.   suffix = s[best[1] + 1] + s[last_idxs[best[1] + 1][0]]
  211.  
  212.   return prefix + middle + suffix
  213.  
  214.  
  215. def brute_force(s, k):
  216.   best = s + "z"
  217.   stack = [(s, k)]
  218.  
  219.   while stack:
  220.     _s, _k = stack.pop()
  221.     if _k == 0:
  222.       best = min(best, _s)
  223.       continue
  224.     stack.append((_s[1:], _k - 1))
  225.     stack.append((_s[0] + _s[2:], _k - 1))
  226.     stack.append((_s[0:len(_s)-1], _k - 1))
  227.     stack.append((_s[0:len(_s)-2] + _s[-1], _k - 1))
  228.  
  229.   return best
  230.  
  231. #    01234567
  232. #s = "abacaaba"
  233. #k = 2
  234.  
  235.  
  236. def lex_kr(x,K,k_r):
  237.     """
  238.     Get a lexicographically comparable subset of `x` for a given value of
  239.     `k_r`.
  240.     """
  241.     N = len(x)
  242.     assert k_r > 0 and k_r < K # check for corner cases
  243.     k_l = K - k_r
  244.     v_l = min(x[:k_l+1])
  245.     v_r = min(x[-k_r-1:])
  246.  
  247.     lex = [v_l]
  248.     lex += x[k_l+1:N-k_r-1]
  249.     lex += [v_r]
  250.  
  251.     return lex
  252.  
  253. def lex_corner(x,K):
  254.     """
  255.     Get the two lexicographically comparable subsets of `x` for corner cases
  256.     when `k_r=0` and `k_r=K`.
  257.     """
  258.     N = len(x)
  259.  
  260.     k_l = K
  261.     v_l = min(x[:k_l+1])
  262.     lex0 = [v_l]
  263.     lex0 += x[k_l+1:]
  264.  
  265.     k_r = K
  266.     v_r = min(x[-k_r-1:])
  267.     lex1 = x[:N-k_r-1]
  268.     lex1 += [v_r]
  269.  
  270.     return lex0,lex1
  271.  
  272.  
  273. def min_lex(x,K):
  274.     subsets = [ lex_kr(x,K,k_r) for k_r in range(1,K) ]
  275.     subsets += lex_corner(x,K) # append the two corner cases
  276.     return min(subsets)
  277.  
  278.  
  279.  
  280. # Test
  281. import random
  282.  
  283. n = 10
  284. num_tests = 100
  285.  
  286. for _ in range(num_tests):
  287.   s = "".join([chr(97 + random.randint(0, 25)) for i in range(n)])
  288.   k = random.randint(1, n - 5)
  289.   #print(s, k)
  290.   _f = [x for x in f(s, k)]
  291.   #rute = brute_force(s, k)
  292.   brute = min_lex([x for x in s], k)
  293.   if brute != _f:
  294.     print("MISMATCH!")
  295.     print(s, k)
  296.     print(_f)
  297.     print(brute)
  298.     break
  299.  
  300. print("Done.")
Success #stdin #stdout 0.03s 11624KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
Done.
https://ideone.com/NNXid3
language:
Python 3 (python  3.12)
created:
4 years ago
visibility:
public

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!