Use this to learn the idea, then write your own version.
123 45class Solution(object):6 def smallestTrimmedNumbers(self, nums, queries):7 """8 :type nums: List[str]9 :type queries: List[List[int]]10 :rtype: List[int]11 """12 max_t = max(t for _, t in queries)13 lookup = [[] for _ in xrange(max_t+1)]14 for i, (k, t) in enumerate(queries):15 lookup[t].append((k, i))16 result = [0]*len(queries)17 idxs = range(len(nums))18 for l in xrange(1, max_t+1):19 cnt = [0]*1020 for i in idxs:21 d = int(nums[i][-l])22 cnt[d] += 123 for d in xrange(9):24 cnt[d+1] += cnt[d]25 new_idxs = [0]*len(nums)26 for i in reversed(idxs):27 d = int(nums[i][-l])28 cnt[d] -= 129 new_idxs[cnt[d]] = i30 idxs = new_idxs31 for k, i in lookup[l]:32 result[i] = idxs[k-1]33 return result34 35 363738import random39 40 4142class Solution2(object):43 def smallestTrimmedNumbers(self, nums, queries):44 """45 :type nums: List[str]46 :type queries: List[List[int]]47 :rtype: List[int]48 """49 def nth_element(nums, n, compare=lambda a, b: a < b):50 def tri_partition(nums, left, right, target, compare):51 mid = left52 while mid <= right:53 if nums[mid] == target:54 mid += 155 elif compare(nums[mid], target):56 nums[left], nums[mid] = nums[mid], nums[left]57 left += 158 mid += 159 else:60 nums[mid], nums[right] = nums[right], nums[mid]61 right -= 162 return left, right63 64 left, right = 0, len(nums)-165 while left <= right:66 pivot_idx = random.randint(left, right)67 pivot_left, pivot_right = tri_partition(nums, left, right, nums[pivot_idx], compare)68 if pivot_left <= n <= pivot_right:69 return70 elif pivot_left > n:71 right = pivot_left-172 else: 73 left = pivot_right+174 75 def compare(a, b):76 for i in xrange(len(nums[a])-t, len(nums[a])):77 if nums[a][i] < nums[b][i]:78 return True79 if nums[a][i] > nums[b][i]:80 return False81 return cmp(a, b) < 082 83 result = []84 idxs = range(len(nums))85 for k, t in queries:86 nth_element(idxs, k-1, compare=compare)87 result.append(idxs[k-1])88 return result89 90 91929394class Solution3(object):95 def smallestTrimmedNumbers(self, nums, queries):96 """97 :type nums: List[str]98 :type queries: List[List[int]]99 :rtype: List[int]100 """101 def compare(a, b):102 for i in xrange(len(nums[a])-t, len(nums[a])):103 if nums[a][i] < nums[b][i]:104 return -1105 if nums[a][i] > nums[b][i]:106 return 1107 return cmp(a, b)108 109 max_t = max(t for _, t in queries)110 lookup = [[] for _ in xrange(max_t+1)]111 for i, (k, t) in enumerate(queries):112 lookup[t].append((k, i))113 result = [0]*len(queries)114 idxs = range(len(nums))115 for t in xrange(1, max_t+1):116 if not lookup[t]:117 continue118 idxs.sort(cmp=compare)119 for k, i in lookup[t]:120 result[i] = idxs[k-1]121 return result122