Company questionsGroundworkPhonePe

Coding · Union-find

Most Stones Removed with Same Row or Column

  • Appears on company lists
  • Medium
  • 1 source

Company-tagged lists or prep guides name it; no candidate write-up was found.

Solve it in the playground, with tests

The problem
Stones sit on distinct integer grid points, given as [row, col] pairs. A stone can be removed if another stone that has not been removed yet shares its row or its column. Return the largest number of stones that can be removed; for [[0,0],[0,2],[1,1],[2,0],[2,2]] the answer is 3. There can be 1,000 stones with coordinates up to 10,000.

Restated in our own words, not a quote from PhonePe or any page. The answer below is written for this site.

Try it first.The answer is covered until you pick. Talk it through against a 30 minute clock, or just read it.

Clarify first

Pin down the input before you touch the keyboard.

I would ask: are the points distinct? Does sharing a row count even if other stones sit between the two (yes, I assume so)? What are the coordinate ranges, and can the list be empty? Do I return only the count, or also a removal order? Do removals change which stones remain available to the others (yes, only the remaining stones count)?

I would assume distinct points, non-negative coordinates up to 10,000, that two stones are linked whenever they share a row or a column, and that I return just the count. An empty list returns 0.

Brute force

Say the obvious answer and why it is too slow.

Search over removal orders. Pick any stone that currently shares a row or column with another remaining stone, remove it, and recurse on what is left. Track the best count over all branches. Using a bitmask of removed stones as a memo key cuts the repeats, but there are still up to 2^n masks.

That is fine for 10 or 15 stones. At n = 1,000 the number of states is astronomically beyond reach, so enumerating orders cannot work. The insight must come from understanding the structure of which stones can be removed, not from trying orders.

The insight

The one idea that makes it fast, and why it is correct.

Brute force
O(2^n · n²)
Time
O(n α(n)), near linear
Space
O(n)

Build a graph where two stones are joined if they share a row or a column. Stones in the same connected component can all be removed except one, and stones in different components do not interact at all. So the answer is n minus the number of components.

Why a component of k stones loses exactly k - 1: pick a spanning tree of the component. A leaf of that tree shares a row or column with its parent, which is still present, so the leaf can be removed. Remove leaves one after another and only the root is left, so k - 1 removals are possible. The last remaining stone can never go, because any stone it shared a row or column with was in its component and has been removed, and no stone outside the component shares a line with it. So k - 1 is also the maximum.

Comparing every pair of stones to build that graph would cost O(n²). The trick is to make rows and columns the nodes of a union-find structure: a stone at (r, c) simply unions row r with column c. Two stones are in the same component exactly when their row nodes end up in the same set. Rows and columns share the same numbers, so I give them separate names with a prefix, "r" for rows and "c" for columns, so row 3 never merges with column 3. A Map creates entries on demand. With path compression and union by size, each operation is almost constant on average. At the end I count the distinct roots among the stones' row nodes.

Walk it through

A dry run on the example, step by step.

Stones in input order. Each stone unions its row node with its column node (r = row, c = column).
StoneUnionWhat happensComponents
(0,0)r0 + c0new set1
(0,2)r0 + c2c2 joins r0's set1
(1,1)r1 + c1new set2
(2,0)r2 + c0r2 joins c0's set2
(2,2)r2 + c2already same set2

Take [[0,0],[0,2],[1,1],[2,0],[2,2]]. The first stone (0,0) unions row 0 with column 0, which makes one new set. The second stone (0,2) unions row 0 with column 2; row 0 already has a set, so column 2 just joins it. The third stone (1,1) unions row 1 with column 1, a brand new set, so there are now 2 components.

The fourth stone (2,0) unions row 2 with column 0; column 0 is in the first set, so row 2 joins it. The fifth stone (2,2) unions row 2 with column 2, and both are already in the same set, so nothing changes. The table shows 2 components after all five stones, so the answer is 5 - 2 = 3.

The code

Runnable JavaScript, checked against test cases.

removeStones.js
function removeStones(stones) {
  const parent = new Map();
  const size = new Map();
  function find(key) {
    if (!parent.has(key)) {
      parent.set(key, key);
      size.set(key, 1);
    }
    let root = key;
    while (parent.get(root) !== root) root = parent.get(root);
    while (parent.get(key) !== root) {
      const next = parent.get(key);
      parent.set(key, root);
      key = next;
    }
    return root;
  }
  function union(a, b) {
    let rootA = find(a);
    let rootB = find(b);
    if (rootA === rootB) return;
    if (size.get(rootA) < size.get(rootB)) [rootA, rootB] = [rootB, rootA];
    parent.set(rootB, rootA);
    size.set(rootA, size.get(rootA) + size.get(rootB));
  }
  const rowKeys = new Set();
  for (const [row, col] of stones) {
    union("r" + row, "c" + col);
    rowKeys.add("r" + row);
  }
  const components = new Set();
  for (const key of rowKeys) components.add(find(key));
  return stones.length - components.size;
}

Runs here in your browser, in a sandboxed worker. Nothing is sent anywhere.

Try your own input

Write JavaScript that calls the code above and prints with console.log. It runs after the solution.

Edge cases

What to test before you say you are done.

An empty list has no stones and no components, so it returns 0. A single stone is its own component, so it returns 0. Stones that share nothing, such as (1,2) and (2,1), give 0. Here the prefixes matter: without them, row 1 and column 1 would collide and falsely join the two stones.

All stones in one row give n - 1, since the row node unites every column. The code counts components by looking up the root of each stone's row node, so every component is counted once through a stone it actually contains and unused rows or columns are never counted.

The find routine compresses paths with a loop, not recursion, so a long chain cannot overflow the call stack.

Follow-ups to expect

Where the interviewer pushes next.

  1. How would you solve it with a depth-first search over stones that share a line, and what would it cost?
  2. Why can the removal order be ignored when computing the count?
  3. How would you output one valid removal order, not just the count?
  4. If stones arrive one at a time, how would you report the answer after each arrival?
  5. Why use prefixed string keys instead of adding an offset to the column number?

Where candidates slip

The mistakes this problem punishes.

Using the same id space for rows and columns, so row 3 and column 3 become one node and unrelated stones merge. Counting components wrongly, for example returning the component count itself instead of n minus the count.

Other slips: building the pairwise stone graph in O(n²) and then running DFS, which is correct but does more work than the near-linear union of rows and columns; trying to simulate removals by hand and getting stuck on order; forgetting both path compression and union by size, so chains grow long; and returning the number of successful unions, which is the number of row and column nodes minus the components, not the stones minus the components.

Have a go before you read.

Interviews reward the sketch you make cold. Clarify out loud, draw it, then compare with the answer.

Now make it stick

Where this was reported

Public pages only. Companies do not confirm these, and they rotate their questions, so treat each link as one person's account or one list's tag.

Learn the pattern first