Search in Rotated Sorted Array

Rotateされた配列からTargetの数値を探す問題。ソートされた部分を検索範囲としてBinary Searchを行うことでO(log(N))で解くことができる。

pub fn search(nums: Vec<i32>, target: i32) -> i32 {
    let (mut start, mut end) = (0, nums.len() - 1);

    while start <= end {
        let mid = (start + end) / 2;

        if nums[mid] == target {
            return mid as i32;
        }

        if nums[start] <= nums[mid] {
            // if start..=mid is sorted
            if target >= nums[start] && target < nums[mid] {
                // if target is between start and mid, search on left partition
                end = mid - 1;
            } else {
                // else search on right partition
                start = mid + 1;
            }
        } else {
            // if mid..=end is sorted
            if target <= nums[end] && target > nums[mid] {
                // if target is between mid and start, search on right partition
                start = mid + 1;
            } else {
                // else search on left partition
                end = mid - 1;
            }
        }
    }

    -1
}

Median of Two Sorted Arrays

二つのソートされた配列から中央値を探す問題。短い方の配列をベースにBinary Searchを行うことで計算することができる。

use std::cmp::{max, min};

pub fn find_median_sorted_arrays(nums1: Vec<i32>, nums2: Vec<i32>) -> f64 {
    // Set shorter array to nums1
    let (nums1, nums2) = if nums1.len() < nums2.len() {
        (nums1, nums2)
    } else {
        (nums2, nums1)
    };

    let (m, n) = (nums1.len(), nums2.len());

    // initialize i to 0..m
    let (mut i_min, mut i_max) = (0, m);

    // variable used to calculate index j for nums2
    let mid = (m + n + 1) / 2;

    loop {
        let i = (i_min + i_max) / 2;
        let j = mid - i;

        // Run binary search to find median
        if i < i_max && nums2[j - 1] > nums1[i] {
            // i is too small. set range to i+1..i_max
            i_min = i + 1;
        } else if i > i_min && nums1[i - 1] > nums2[j] {
            // i is too big. set range to i_min..i-1
            i_max = i - 1;
        } else {
            // Median is found

            let max_left = if i == 0 {
                // if left partition for nums1 is empty
                nums2[j - 1]
            } else if j == 0 {
                // if left partition for nums2 is empty
                nums1[i - 1]
            } else {
                max(nums1[i - 1], nums2[j - 1])
            };

            // if total length is odd, return the maximum value in left partition
            if (m + n) % 2 == 1 {
                return max_left as f64;
            }

            let min_right = if i == m {
                // if right partition for nums1 is empty
                nums2[j]
            } else if j == n {
                // if right partition for nums2 is empty
                nums1[i]
            } else {
                min(nums1[i], nums2[j])
            };

            return (max_left + min_right) as f64 / 2.0;
        }
    }
}

計算量はO(log(min(m,n)))。

Cut Off Trees for Golf Event

各ノードを始点に各ノードの検査を行うため計算量はO((R x C) ^ 2)。空間計算量はBFSのキューの長さO(R x C)。

use std::collections::{HashSet, VecDeque};

pub fn cut_off_tree(forest: Vec<Vec<i32>>) -> i32 {
    let mut trees = vec![];

    for r in 0..forest.len() {
        for c in 0..forest[r].len() {
            let v = forest[r][c];
            if v > 1 {
                trees.push((v, r, c));
            }
        }
    }

    // Sort by height
    trees.sort();

    let mut ans = 0;
    let mut s_r = 0;
    let mut s_c = 0;

    for (_v, r, c) in trees.into_iter() {
        let d = dist(&forest, s_r, s_c, r, c);
        if d < 0 {
            return -1;
        }
        ans += d;
        s_r = r;
        s_c = c;
    }

    ans
}

fn dist(forest: &[Vec<i32>], s_r: usize, s_c: usize, t_r: usize, t_c: usize) -> i32 {
    let mut queue = VecDeque::new();
    queue.push_back((s_r, s_c, 0));
    let mut seen = HashSet::new();
    seen.insert((s_r, s_c));

    while let Some((r, c, d)) = queue.pop_front() {
        if r == t_r && c == t_c {
            return d;
        }

        for &(d_r, d_c) in [(1, 0), (!0, 0), (0, 1), (0, !0)].iter() {
            let new_r = r.wrapping_add(d_r);
            let new_c = c.wrapping_add(d_c);

            if new_r < forest.len()
                && new_c < forest[r].len()
                && !seen.contains(&(new_r, new_c))
                && forest[new_r][new_c] != 0
            {
                seen.insert((new_r, new_c));
                queue.push_back((new_r, new_c, d + 1));
            }
        }
    }

    -1
}

Diameter of Binary Tree

与えられたBinary Treeから最長のパスを探す問題。DFSで解くとシンプル

use std::cell::RefCell;
use std::cmp::max;
use std::rc::Rc;

#[derive(Debug, PartialEq, Eq)]
pub struct TreeNode {
    pub val: i32,
    pub left: Option<Rc<RefCell<TreeNode>>>,
    pub right: Option<Rc<RefCell<TreeNode>>>,
}

impl TreeNode {
    #[inline]
    pub fn new(val: i32) -> Self {
        TreeNode {
            val,
            left: None,
            right: None,
        }
    }
}

pub fn diameter_of_binary_tree(root: Option<Rc<RefCell<TreeNode>>>) -> i32 {
    let mut ans = 0;

    depth(root.clone(), &mut ans);

    ans
}

fn depth(node: Option<Rc<RefCell<TreeNode>>>, ans: &mut i32) -> i32 {
    if let Some(node) = node.as_ref() {
        let left = depth(node.borrow().left.clone(), ans);
        let right = depth(node.borrow().right.clone(), ans);
        *ans = max(*ans, left + right);
        max(left, right) + 1
    } else {
        0
    }
}

計算量はO(N)

Lowest Common Ancestor of a Binary Tree

与えられたBinary Treeから共通の先祖を探す問題。DFSで解くことができる。

use std::cell::RefCell;
use std::rc::Rc;

#[derive(Debug, PartialEq, Eq)]
pub struct TreeNode {
    pub val: i32,
    pub left: Option<Rc<RefCell<TreeNode>>>,
    pub right: Option<Rc<RefCell<TreeNode>>>,
}

impl TreeNode {
    #[inline]
    pub fn new(val: i32) -> Self {
        TreeNode {
            val,
            left: None,
            right: None,
        }
    }
}
pub fn lowest_common_ancestor(
    root: Option<Rc<RefCell<TreeNode>>>,
    p: Option<Rc<RefCell<TreeNode>>>,
    q: Option<Rc<RefCell<TreeNode>>>,
) -> Option<Rc<RefCell<TreeNode>>> {
    let root = root.unwrap();
    let p = p.unwrap();
    let q = q.unwrap();
    let mut ans = root.clone();

    recurse(root.clone(), p.clone(), q.clone(), &mut ans);

    Some(ans)
}

fn recurse(
    node: Rc<RefCell<TreeNode>>,
    p: Rc<RefCell<TreeNode>>,
    q: Rc<RefCell<TreeNode>>,
    ans: &mut Rc<RefCell<TreeNode>>,
) -> bool {
    let left = if let Some(left) = node.borrow().left.as_ref() {
        recurse(left.clone(), p.clone(), q.clone(), ans)
    } else {
        false
    };

    let right = if let Some(right) = node.borrow().right.as_ref() {
        recurse(right.clone(), p.clone(), q.clone(), ans)
    } else {
        false
    };

    let mid = node == p || node == q;

    if mid as i32 + left as i32 + right as i32 >= 2 {
        *ans = node.clone();
    }

    mid || left || right
}

時間計算量、空間計算量の両方がノードの数O(N)。

Course Schedule

グラフの中からサイクルを検出する問題。DFSでO(E + V)で解くことができる。

use std::collections::HashMap;

pub fn can_finish(num_courses: i32, prerequisites: Vec<Vec<i32>>) -> bool {
    let mut courses = HashMap::new();
    let num_courses = num_courses as usize;

    for relation in prerequisites.into_iter() {
        let (next, prerequisite) = (relation[0] as usize, relation[1] as usize);
        courses.entry(prerequisite).or_insert(Vec::new()).push(next);
    }

    let mut checked = vec![false; num_courses];
    let mut path = vec![false; num_courses];

    for i in 0..num_courses {
        if is_cycle(i, &courses, &mut checked, &mut path) {
            return false;
        }
    }

    true
}

fn is_cycle(
    curr_course: usize,
    courses: &HashMap<usize, Vec<usize>>,
    checked: &mut [bool],
    path: &mut [bool],
) -> bool {
    // It will not be a cycle if node is previously checked
    if checked[curr_course] {
        return false;
    }

    // If node is visited in current path, it is a cycle
    if path[curr_course] {
        return true;
    }

    path[curr_course] = true;

    if let Some(children) = courses.get(&curr_course) {
        for &child in children.iter() {
            if is_cycle(child, courses, checked, path) {
                return true;
            }
        }
    }

    path[curr_course] = false;
    checked[curr_course] = true;

    false
}

Number of Islands

与えられた二次元の配列から上下左右で繋がっている島の数を数える問題。順番に各IndexをIterateし、'1'に遭遇した際にCountをIncrementし、DFSで上下左右の'1'を'0'に書き換えることで解くことができる。

pub fn num_islands(grid: Vec<Vec<char>>) -> i32 {
    let mut count = 0;
    let mut grid = grid;

    for i in 0..grid.len() {
        for j in 0..grid[i].len() {
            if grid[i][j] == '1' {
                count += 1;
                dfs(&mut grid, i, j);
            }
        }
    }

    count
}

fn dfs(grid: &mut [Vec<char>], i: usize, j: usize) {
    grid[i][j] = '0';
    for &(d_i, d_j) in [(1, 0), (!0, 0), (0, 1), (0, !0)].iter() {
        let new_i = i.wrapping_add(d_i);
        let new_j = j.wrapping_add(d_j);
        if new_i < grid.len() && new_j < grid[new_i].len() && grid[new_i][new_j] == '1' {
            dfs(grid, new_i, new_j);
        }
    }
}

時間計算量は各ノードを一度訪問するのでO(M x N)、空間計算量はスタックがDFSで最大でもO(M x N)。