遇到了挺多有意思的题目，记录一下。可能包含一些 CUDA Kernel。

LRU Cache

当时面试官说KV Cache，想了想说实现 LRU 功能就行，那就直接 leetcode 原题了。

CUDA Reduce

经典的 reduce kernel。写到了 warp shuffle。

__global__ reduce_kernel(const float* input, float* output, int n) {
    float val = 0.0f;

    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i += blockDim.x * gridDim.x) {
        val += input[i];
    }

    __shared__ float shared[TILE];

    // warp reduce
    for (int offset = warpSize / 2; offset > 0; offset /= 2) {
        val += shfl_down_sync(0xffffffff, val, offset);
    }

    int lane = threadIdx.x & 31;
    int wid = threadIdx / 32;

    if (lane == 0) {
        shared[wid] = val;
    }
    __syncthreads();

    if (threadIdx.x < (blockDim.x + 31) / 32) {
        val = shared[lane];
    }
    else {
        val = 0.0f;
    }

    if (wid == 0) {
        for (int offset = warpSize / 2; offset > 0; offset /= 2) {
            val += shfl_down_sync(0xffffffff, val, offset);
        }
    }

    if (threadIdx.x == 0) {
        output[blockIdx.x] = val;
    }
}

数组截断

有若干个int数组，要按相同比例x%截断，每个数组保留前x%的元素（向下取整，x可为浮点数）。现要求最终留下来的数去重后不多于n个，问最多能留下多少。 例子：

• 假设 arrs = {{1, 2, 3, 4}, {2, 5, 1, 4, 3}} n = 4 那截断比例取75%时 arrs = {{1, 2, 3}, {2, 5, 1}} 去重后有4个，分别是{1, 2, 3, 5} 继续增加元素就会超过n，因此最终输出 {{1, 2, 3}, {2, 5, 1}}

void clip(std::vector<std::vector<int>> &arrs, int n)

这题面试的时候想的方法是对比例做二分，当截断的矩阵长度不再变化时退出。

// arr[:x%], s.t. non-duplicate numbers <= n
void clip(std::vector<std::vector<int>> &arrs, int n) {
    float p_left = 0.0, p_right = 1.0;
    // bool flag = true;

    std::vector<int> clipped_lengths_prev;
    while (true) {
        float p = (p_left + p_right) / 2;
    
        std::unordered_set<int> nums;

        for (auto arr : arrs) {
            std::vector<int> clipped_lengths_curr;
            int n = arr.size();
            int clipped_n = floor(n * p);
            clipped_lengths_curr.push_back(clipped_n);
            for (int i = 0; i < clipped_n; i++) {
                auto it = nums.find(arr[i]);
                if (it == nums.end()) {
                    nums.insert(arr[i]);
                }
            }
        }
        if (nums.size() > n) { // 不满足条件
            // int num_arr = arrs.size();
            // for (int i = 0; i < num_arr; i++) {
            //     arrs[i] = arrs[i][:clipped_lengths_prev[i]];
            // }
            p_right = p;
        }
        else if (nums.size() == n) {
            int num_arr = arrs.size();
            for (int i = 0; i < num_arr; i++) {
                arrs[i] = arrs[i][:clipped_lengths_curr[i]];
            }
            return;
        }
        else {
            p_left = p;
        }
        if (clipped_lengths_prev == clipped_lengths_curr) { // 已收敛
            int num_arr = arrs.size();
            for (int i = 0; i < num_arr; i++) {
                arrs[i] = arrs[i][:clipped_lengths_prev[i]];
            }
            return;
        }
        clipped_lengths_prev = clipped_lengths_curr;
    }
    return;
}

最大全1矩形面积

给定仅含0、1的矩阵，找出只含1的最大矩形，输出其面积。

一开始想的是经典的动态规划，但DP的条件是正方形，而这题是矩形。正确的做法是逐行维护一个 heights 数组，heights[i] 代表第 $i$ 列地当前元素及其上方有几个连续的 1。举个例子，对于矩阵 A = [[1, 0, 0, 1], [1, 1, 1, 0]]，维护到第二行的时候，heights 应为 [2,1,1,0]。对于每一行，就相当于在求解 leetcode 经典题目之盛最多水的容器。

假设矩阵大小 $M\times N$，则不难分析得到时间复杂度为 $O(MN)$

int solve_per_row(const vector<int>& nums) {
    int n = nums.size();
    int left = 0; right = n - 1;
    int res = 0;
    while (left < right) {
        res = max(res, min(nums[left], nums[right]) * (right - left));
        if (nums[left] < nums[right]) {
            left--;
        }
        else {
            right++;
        }
    }
    return res;
}

int solve(const vector<vector<int>>& matrix) {
    int m = matrix.size(), n = matrix[0].size();
    vector<int> heights(n, 0);
    int res = 0;
    for (int i = 0; i < m; i++) {
        // Update heights
        for (int j = 0; j < n; j++) {
            if (i == 0) {
                heights[j] = matrix[0][j];
            }
            else {
                heights[j] = (matrix[i][j] == 0) ? 0 : heights[j] + 1;
            }
        }
        
        res = max(res, solve_per_row(heights));
    }
    return res;
}

最大乘积

给定N个整数和一个K，选出其中K个整数，使乘积尽可能大。要求返回这 K 个数。

做法是贪心。先排序，然后从两端成对地选。时间复杂度 $O(n\log n)$

vector<int> k_nums_for_max_product(const vector<int>& nums, int K) {
    int N = nums.size();
    vector<int> res;
    sort(nums.begin(), nums.end());

    if ((K & 1) && nums[N - 1] <= 0) {
        for (int i = 0; i < K; i++) {
            res.push_back(nums[N - 1 - i]);
        }
        return res;
    }

    int remain = K;
    int left = 0, right = N - 1;
    if (remain & 1) { // K 为奇数，先挑选最大的正数
        res.push_back(nums[right]);
        right--;
        remain--;
    }
    while (remain > 0) {
        int left_product = nums[left] * nums[left + 1]; // negative couple
        int right_product = nums[right] * nums[right - 1];
        if (left_product > right_product) {
            res.push_back(nums[left]);
            res.push_back(nums[left + 1]);
            left += 2;
        }
        else {
            res.push_back(nums[right]);
            res.push_back(nums[right - 1]);
            right -= 2;
        }
        remain -= 2;
    }

    return res;
}

实际上正确的解法应该整点离散的，枚举出所有切分的点然后统计。

手撕 DataLoader

没错，手写一个 PyTorch 的 DataLoader。是一家量化公司的题目。当时就写了一个 __iter__ 方法和 batch 队列预处理，就不把我的代码放出来献丑了。

排队掷骰

N个人排队掷骰子拿一个物品，如果点数为6，则拿到，点数为1，则淘汰，否则挪到队尾排队，求第i个人最终拿到的概率。

概率 dp，在面试官的提示下写出了转移方程。

首先是定义。记 $P(n, i)$ 为队列长度为 $n$ 时，第 $i$ 个人最终获物品的概率。那么我们可以得出

• $P(n, 1) = \frac{1}{6} + \frac{2}{3}P(n, n)$，即第一个人要获得奖品，要么直接扔出6，要么扔出2到5，然后成为队伍的第 $n$ 个人。
• $P(n, i) = \frac{1}{6}P(n - 1, i - 1) + \frac{2}{3}P(n, i - 1), i > 1$，即对于第 $i (i>1)$ 个人，需要当前队首的人要么扔到1被淘汰，队列长度减1，自己变成第 $i - 1$ 个人；或者队首的人扔到2～5，去队尾，队列长度仍为 $n$ ，而自己变成第 $i - 1$ 个人。

我们需要的就是从这两个式子中找到递推公式。首先我们需要一个边界条件。代入 $n = 1$ 计算 $P(n, 1)$，不难得到 $P(1, 1)=\frac{1}{2}$。对于每个 $n$，我们计算出 $P(n, 1)$ 后即可递推计算得到后续的 $P(n, i)$。经过一通等比数列的计算，可以得到

$P(n, 1)=\frac{\frac{1}{6} + \frac{1}{9}\sum_{k=1}^{n-1} (\frac{2}{3})^{n - 1 - k}P(n - 1, k)}{1-(\frac{2}{3})^n}$

然后用 $P(n, i)$ 递推就可以了。

最长公共前缀

给定 n 个字符串，寻找一个最长字符串，使它至少是其中 k 个字符串的前缀。解决方法是前缀树(Trie)

#include <iostream>
#include <string>
#include <vector>
#include <unordered_map>
using namespace std;

class Trie {
private:
    struct Node {
        unordered_map<char, int> next;
        int count = 0;  // 有多少个字符串经过该节点
    };

    vector<Node> tree;
    int k;
    string answer;

    void dfs(int node, string& path) {
        if (tree[node].count < k) {
            return;
        }

        if (path.size() > answer.size()) {
            answer = path;
        }

        for (const auto& [ch, nextNode] : tree[node].next) {
            path.push_back(ch);
            dfs(nextNode, path);
            path.pop_back();
        }
    }

public:
    Trie(int k) : k(k) {
        tree.emplace_back();  // 根节点
    }

    void insert(const string& s) {
        int node = 0;

        for (char ch : s) {
            if (!tree[node].next.count(ch)) {
                tree[node].next[ch] = tree.size();
                tree.emplace_back();
            }

            node = tree[node].next[ch];
            ++tree[node].count;
        }
    }

    string getLongestPrefix() {
        string path;
        dfs(0, path);
        return answer;
    }
};

string longestCommonPrefixAtLeastK(const vector<string>& sequences, int k) {
    if (k <= 0 || k > static_cast<int>(sequences.size())) {
        return "";
    }

    Trie trie(k);

    for (const string& sequence : sequences) {
        trie.insert(sequence);
    }

    return trie.getLongestPrefix();
}

int main() {
    vector<string> sequences = {
        "flower",
        "flow",
        "flight",
        "flock",
        "dog"
    };

    int k = 3;

    string result = longestCommonPrefixAtLeastK(sequences, k);
    cout << result << '\n';  // 输出：flo

    return 0;
}

Ring All-Reduce

用CUDA实现 Ring All-reduce。其实写个伪代码就差不多了，体现 reduce-scatter 和 all gather 两个阶段就可以。