LeetCode 2343 - Query Kth Smallest Trimmed Number

This problem provides a list of numbers represented as strings, all of equal length. The task is to answer multiple queries, where each query asks for the index of the k-th smallest number after trimming every number in the list to its last trimi digits.

LeetCode Problem 2343

Difficulty: 🟡 Medium
Topics: Array, String, Divide and Conquer, Sorting, Heap (Priority Queue), Radix Sort, Quickselect

Solution

Problem Understanding

This problem provides a list of numbers represented as strings, all of equal length. The task is to answer multiple queries, where each query asks for the index of the k-th smallest number after trimming every number in the list to its last trimi digits. Importantly, if two trimmed numbers are equal, the one with the lower original index is considered smaller. After each query, the numbers must be reset to their original form for the next query.

The input consists of a list nums of strings containing only digits, with lengths up to 100, and a list of queries, each with a k value and a trim length. The expected output is a list of indices corresponding to each query result.

The key points are: trimming from the right, comparison of numbers lexicographically, maintaining original indices for tie-breaking, and efficiently handling multiple queries. Potential edge cases include numbers with leading zeros, queries that request the full number without trimming, and multiple identical trimmed numbers.

Approaches

The brute-force approach is straightforward: for each query, trim every number in the list, sort the trimmed numbers while keeping track of original indices, and select the k-th smallest number. This approach works but can be inefficient because it requires sorting for each query independently.

The optimal approach leverages preprocessing by trim length. We can store trimmed numbers for each possible trim length and sort them once, then answer each query in constant time using the precomputed results. This reduces redundant computation and allows us to handle multiple queries efficiently.

Approach Time Complexity Space Complexity Notes
Brute Force O(Q * N log N) O(N) Trim and sort for each query independently
Optimal O(N * L + Q) O(N * L) Precompute sorted indices for all trim lengths, then answer queries quickly

Here, N is the number of numbers, L is the length of each number, and Q is the number of queries.

Algorithm Walkthrough

  1. Determine the maximum possible trim length from the length of the numbers.
  2. Precompute trimmed numbers for every possible trim length from 1 to the full length.
  3. For each trim length, pair each trimmed number with its original index and sort the list lexicographically, breaking ties by index.
  4. Store only the sorted indices for each trim length in a dictionary or array for fast lookup.
  5. For each query [k, trim], retrieve the precomputed sorted indices for the given trim length and select the (k-1)-th index as the answer.
  6. Collect all query results into the final answer list and return it.

Why it works: Sorting trimmed numbers with their original indices ensures that ties are broken correctly. Precomputing for all trim lengths avoids redundant computation for queries with the same trim length, which guarantees both correctness and efficiency.

Python Solution

from typing import List

class Solution:
    def smallestTrimmedNumbers(self, nums: List[str], queries: List[List[int]]) -> List[int]:
        n = len(nums)
        l = len(nums[0])
        trimmed_indices = {}
        
        # Precompute sorted indices for each trim length
        for trim in range(1, l + 1):
            trimmed_list = [(num[-trim:], idx) for idx, num in enumerate(nums)]
            trimmed_list.sort()  # Sorts by string first, then index due to tuple ordering
            trimmed_indices[trim] = [idx for _, idx in trimmed_list]
        
        answer = []
        for k, trim in queries:
            answer.append(trimmed_indices[trim][k-1])
        
        return answer

This Python implementation first precomputes sorted indices for every possible trim length. Each trimmed number is paired with its original index, ensuring tie-breaking is correct. After sorting, only the indices are stored for fast query responses. Each query simply looks up the appropriate precomputed list and retrieves the k-th smallest index.

Go Solution

import (
    "sort"
    "strconv"
)

func smallestTrimmedNumbers(nums []string, queries [][]int) []int {
    n := len(nums)
    l := len(nums[0])
    
    trimmedIndices := make(map[int][]int)
    
    // Precompute sorted indices for each trim length
    for trim := 1; trim <= l; trim++ {
        type pair struct {
            val string
            idx int
        }
        trimmedList := make([]pair, n)
        for i, num := range nums {
            trimmedList[i] = pair{num[len(num)-trim:], i}
        }
        sort.Slice(trimmedList, func(i, j int) bool {
            if trimmedList[i].val == trimmedList[j].val {
                return trimmedList[i].idx < trimmedList[j].idx
            }
            return trimmedList[i].val < trimmedList[j].val
        })
        sortedIdx := make([]int, n)
        for i, p := range trimmedList {
            sortedIdx[i] = p.idx
        }
        trimmedIndices[trim] = sortedIdx
    }
    
    answer := make([]int, len(queries))
    for i, q := range queries {
        k, trim := q[0], q[1]
        answer[i] = trimmedIndices[trim][k-1]
    }
    
    return answer
}

In Go, the solution is similar. The primary differences include the use of structs to pair values and indices, and sort.Slice for custom sorting. Maps store precomputed sorted indices by trim length, which are then accessed to efficiently answer queries.

Worked Examples

Example 1

nums = ["102","473","251","814"], queries = [[1,1],[2,3],[4,2],[1,2]]
  1. Trim length 1: ["2","3","1","4"], sorted indices: [2,0,1,3]. Query [1,1] → 2.
  2. Trim length 3: ["102","473","251","814"], sorted indices: [0,2,1,3]. Query [2,3] → 2.
  3. Trim length 2: ["02","73","51","14"], sorted indices: [0,3,2,1]. Query [4,2] → 1.
  4. Trim length 2 again: Query [1,2] → 0.

Example 2

nums = ["24","37","96","04"], queries = [[2,1],[2,2]]
  1. Trim length 1: ["4","7","6","4"], sorted indices: [0,3,2,1]. Query [2,1] → 3.
  2. Trim length 2: ["24","37","96","04"], sorted indices: [3,0,1,2]. Query [2,2] → 0.

Complexity Analysis

Measure Complexity Explanation
Time O(N * L log N + Q) Precompute sorted indices for each trim length (N numbers, up to L trim lengths), sorting each list takes O(N log N). Query lookup is O(1) per query.
Space O(N * L) Storing sorted indices for each trim length requires N indices for each of L possible trims.

This approach is efficient for the problem constraints, as N and L are both ≤ 100, making precomputation feasible.

Test Cases

# Provided examples
assert Solution().smallestTrimmedNumbers(["102","473","251","814"], [[1,1],[2,3],[4,2],[1,2]]) == [2,2,1,0]
assert Solution().smallestTrimmedNumbers(["24","37","96","04"], [[2,1],[2,2]]) == [3,0]

# Edge cases
assert Solution().smallestTrimmedNumbers(["0001","0002","0003"], [[1,1],[3,4]]) == [0,2]  # Leading zeros
assert Solution().smallestTrimmedNumbers(["1"], [[1,1]]) == [0]  # Single element
assert Solution().smallestTrimmedNumbers(["10","01"], [[1,2],[2,1]]) == [1,0]  # Tie-breaking by index
assert Solution().smallestTrimmedNumbers(["123","456","789"], [[3,2]]) == [2]  # Trim shorter than full length
Test Why
["102","473","251","814"] Validates multiple trims and queries
["24","37","96","04"] Checks tie-breaking on last digit
["0001","0002","0003"] Handles leading zeros correctly
["1"] Single element list
["10","01"] Ensures index tie-breaking works
["123","456","789"] Confirms trimming shorter than full length

Edge Cases

Leading zeros: Numbers like "001" must preserve the trimming rules and correct comparison. The implementation uses string comparison, which handles leading zeros correctly.

Single element list: Queries on a single-element array should always return index 0. The solution naturally handles this as sorting one element is trivial.

Tie-breaking by index: When multiple trimmed numbers are identical, the one with the smaller original index should be chosen. By storing pairs of (trimmed_number, index) and sorting tuples, the implementation respects this rule automatically.

**Maximum trim length